"""Snapshot selection and filesystem writers for runtime collection histories.""" from __future__ import annotations import json from pathlib import Path from typing import Any, TypeVar import numpy as np from python_app.models.dataset_model import ResultCollection, SweepCollection from python_app.storage.npz.paths import collection_dir_name, sanitize_path_component from python_app.storage.npz.serialize import serialize_result_collection, serialize_trace_collection TCollection = TypeVar("TCollection") def select_aligned_histories( raw_history: list[SweepCollection], preprocessed_history: list[SweepCollection], result_history: list[ResultCollection], last_n: int, ) -> tuple[list[SweepCollection], list[SweepCollection], list[ResultCollection], dict[str, Any]]: """Select history tails prioritizing currently displayed processed results.""" raw_index, _ = _index_by_collection_sequence(raw_history) pre_index, _ = _index_by_collection_sequence(preprocessed_history) result_index, result_pos = _index_by_collection_sequence(result_history) if result_index: ordered_result_keys = sorted(result_index, key=lambda key: result_pos[key]) selected_keys = ordered_result_keys[-last_n:] return ( [raw_index[key] for key in selected_keys if key in raw_index], [pre_index[key] for key in selected_keys if key in pre_index], [result_index[key] for key in selected_keys], { "selection_mode": "result_tail_with_optional_alignment", "selected_collection_ids": [int(key[0]) for key in selected_keys], }, ) return ( raw_history[-last_n:], preprocessed_history[-last_n:], result_history[-last_n:], { "selection_mode": "independent_tail", "selected_collection_ids": [], }, ) def save_trace_history_binary(stage_dir: Path, history: list[SweepCollection], magic: int) -> None: """Write binary trace history with lightweight metadata sidecars.""" stage_dir.mkdir(parents=True, exist_ok=True) for index, collection in enumerate(history): binary_path = stage_dir / f"{index:04d}.bin" metadata_path = stage_dir / f"{index:04d}.json" binary_path.write_bytes(serialize_trace_collection(collection, magic)) metadata_path.write_text( json.dumps( { "collection_id": collection.collection_id, "monotonic_ns": collection.monotonic_ns, "trace_count": len(collection.traces), }, indent=2, ), encoding="utf-8", ) def save_result_history_binary(stage_dir: Path, history: list[ResultCollection]) -> None: """Write binary processed-result history with metadata sidecars.""" stage_dir.mkdir(parents=True, exist_ok=True) for index, collection in enumerate(history): binary_path = stage_dir / f"{index:04d}.bin" metadata_path = stage_dir / f"{index:04d}.json" binary_path.write_bytes(serialize_result_collection(collection)) metadata_path.write_text( json.dumps( { "collection_id": collection.collection_id, "monotonic_ns": collection.monotonic_ns, "block_count": len(collection.blocks), }, indent=2, ), encoding="utf-8", ) def save_trace_history_numpy(stage_dir: Path, history: list[SweepCollection]) -> None: """Write raw/preprocessed collections as NumPy directory tree.""" stage_dir.mkdir(parents=True, exist_ok=True) for index, collection in enumerate(history): collection_dir = stage_dir / collection_dir_name(index, collection.collection_id, collection.monotonic_ns) collection_dir.mkdir(parents=True, exist_ok=False) traces_meta: list[dict[str, int | str]] = [] for trace in collection.traces: tag = f"i{trace.combo.input_pos}_o{trace.combo.output_pos}" freq = np.asarray(trace.frequency_hz, dtype=np.float32) s21 = np.asarray(trace.s21, dtype=np.complex64) np.save(collection_dir / f"{tag}_freq.npy", freq) np.save(collection_dir / f"{tag}_s21.npy", s21) traces_meta.append( { "input": int(trace.combo.input_pos), "output": int(trace.combo.output_pos), "points": int(freq.size), "freq_file": f"{tag}_freq.npy", "s21_file": f"{tag}_s21.npy", } ) (collection_dir / "meta.json").write_text( json.dumps( { "collection_id": int(collection.collection_id), "monotonic_ns": int(collection.monotonic_ns), "trace_count": len(collection.traces), "traces": traces_meta, }, indent=2, ), encoding="utf-8", ) def save_result_history_numpy(stage_dir: Path, history: list[ResultCollection]) -> None: """Write processed result collections as NumPy directory tree.""" stage_dir.mkdir(parents=True, exist_ok=True) for index, collection in enumerate(history): collection_dir = stage_dir / collection_dir_name(index, collection.collection_id, collection.monotonic_ns) collection_dir.mkdir(parents=True, exist_ok=False) blocks_meta: list[dict[str, int | str | list[dict[str, int | str | float]]]] = [] for block_index, block in enumerate(collection.blocks): block_dir = collection_dir / f"block_{block_index:03d}_i{block.combo.input_pos}_o{block.combo.output_pos}" block_dir.mkdir(parents=True, exist_ok=False) payload_meta: list[dict[str, int | str | float]] = [] for payload_index, payload in enumerate(block.payloads): safe_name = sanitize_path_component(payload.processing_name or "processor") base_name = f"{payload_index:03d}_{safe_name}_kind{payload.kind}" if payload.kind == 1: freq = np.asarray(payload.frequency_hz, dtype=np.float32) trace = np.asarray(payload.trace, dtype=np.complex64) np.save(block_dir / f"{base_name}_freq.npy", freq) np.save(block_dir / f"{base_name}_trace.npy", trace) payload_meta.append( { "kind": int(payload.kind), "name": payload.processing_name, "points": int(freq.size), "freq_file": f"{base_name}_freq.npy", "trace_file": f"{base_name}_trace.npy", } ) elif payload.kind == 2: scalar = np.asarray([float(payload.scalar_value)], dtype=np.float32) np.save(block_dir / f"{base_name}_scalar.npy", scalar) payload_meta.append( { "kind": int(payload.kind), "name": payload.processing_name, "scalar_file": f"{base_name}_scalar.npy", "scalar_value": float(payload.scalar_value), } ) blocks_meta.append( { "input": int(block.combo.input_pos), "output": int(block.combo.output_pos), "payload_count": len(block.payloads), "dir": block_dir.name, "payloads": payload_meta, } ) (collection_dir / "meta.json").write_text( json.dumps( { "collection_id": int(collection.collection_id), "monotonic_ns": int(collection.monotonic_ns), "block_count": len(collection.blocks), "blocks": blocks_meta, }, indent=2, ), encoding="utf-8", ) def validate_snapshot_name(snapshot_name: str) -> str: """Normalize snapshot directory stem.""" return sanitize_path_component(snapshot_name) def _index_by_collection_sequence( history: list[TCollection], ) -> tuple[dict[tuple[int, int], TCollection], dict[tuple[int, int], int]]: """Index collections by latest-first `(collection_id, occurrence)` preserving history position. Occurrence is counted from the tail (newest item is occurrence `0` for its collection id). This avoids cross-run misalignment when collection ids are reused after restarts and one stage keeps longer history than another. """ indexed: dict[tuple[int, int], TCollection] = {} positions: dict[tuple[int, int], int] = {} seen_count_from_tail_by_id: dict[int, int] = {} # Walk from newest to oldest so occurrence=0 always means the most recent # instance for the given collection id. for index in range(len(history) - 1, -1, -1): item = history[index] collection_id = int(getattr(item, "collection_id")) occurrence = seen_count_from_tail_by_id.get(collection_id, 0) key = (collection_id, occurrence) indexed[key] = item positions[key] = index seen_count_from_tail_by_id[collection_id] = occurrence + 1 return indexed, positions