"""Snapshot selection and filesystem writers for runtime collection histories.""" from __future__ import annotations import json import logging from contextlib import contextmanager from pathlib import Path import shutil import tempfile from typing import Any, Iterator import numpy as np from python_app.models.dataset_model import ResultCollection, ResultPayload, 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, validate_trace_collection logger = logging.getLogger(__name__) @contextmanager def atomic_snapshot_directory(destination: Path) -> Iterator[Path]: """Publish a complete export directory; discard staging files after a failed write.""" destination.parent.mkdir(parents=True, exist_ok=True) if destination.exists(): raise FileExistsError(f"Snapshot directory already exists: {destination}") staging = Path(tempfile.mkdtemp(prefix=f".{destination.name}.", dir=destination.parent)) try: yield staging if destination.exists(): raise FileExistsError(f"Snapshot directory already exists: {destination}") staging.rename(destination) finally: if staging.exists(): shutil.rmtree(staging) def _matching_collection( history: list[SweepCollection], anchor: SweepCollection | ResultCollection, used: set[int], *, same_timestamp: bool, ) -> SweepCollection | None: """Match preprocessing identities exactly and raw sweeps by their capture window. Producers stamp raw/preprocessed collections at different instants. Their capture window is copied unchanged; results keep the preprocessed timestamp. Old files without timing fall back to the latest preceding collection with the same id. """ candidates = [ (index, item) for index, item in enumerate(history) if index not in used and int(item.collection_id) == int(anchor.collection_id) ] timestamp = int(anchor.monotonic_ns) start_ns = int(getattr(anchor, "capture_start_ns", 0)) end_ns = int(getattr(anchor, "capture_end_ns", 0)) if same_timestamp and timestamp: candidates = [(index, item) for index, item in candidates if int(item.monotonic_ns) == timestamp] elif start_ns or end_ns: candidates = [ (index, item) for index, item in candidates if (int(item.capture_start_ns), int(item.capture_end_ns)) == (start_ns, end_ns) ] elif timestamp: candidates = [(index, item) for index, item in candidates if int(item.monotonic_ns) <= timestamp] if not candidates: return None index, item = max(candidates, key=lambda pair: (int(pair[1].monotonic_ns), pair[0])) used.add(index) return item def select_aligned_histories( raw_history: list[SweepCollection], preprocessed_history: list[SweepCollection], result_history: list[ResultCollection], last_n: int, *, include_unaligned: bool = False, ) -> tuple[list[SweepCollection], list[SweepCollection], list[ResultCollection], dict[str, Any]]: """Select aligned tails, or retain every stage independently for complete exports.""" if last_n <= 0: raise ValueError("last_n must be > 0") if include_unaligned: return ( raw_history[-last_n:], preprocessed_history[-last_n:], result_history[-last_n:], {"selection_mode": "independent_tail", "anchor_stage": "none", "selected_collection_ids": [int(item.collection_id) for item in result_history[-last_n:]], "aligned_key_count": 0}, ) selected_raw: list[SweepCollection] = [] selected_pre: list[SweepCollection] = [] selected_results = result_history[-last_n:] used_raw: set[int] = set() used_pre: set[int] = set() if selected_results: anchors = selected_results anchor_stage = "results" for result in reversed(selected_results): pre = _matching_collection(preprocessed_history, result, used_pre, same_timestamp=True) raw = _matching_collection(raw_history, pre or result, used_raw, same_timestamp=False) if pre is not None: selected_pre.append(pre) if raw is not None: selected_raw.append(raw) selected_pre.reverse() selected_raw.reverse() elif preprocessed_history: selected_pre = preprocessed_history[-last_n:] anchors = selected_pre anchor_stage = "preprocessed" for pre in reversed(selected_pre): raw = _matching_collection(raw_history, pre, used_raw, same_timestamp=False) if raw is not None: selected_raw.append(raw) selected_raw.reverse() else: selected_raw = raw_history[-last_n:] anchors = selected_raw anchor_stage = "raw" if anchors else "none" return selected_raw, selected_pre, selected_results, { "selection_mode": f"{anchor_stage}_tail_with_optional_alignment", "anchor_stage": anchor_stage, "selected_collection_ids": [int(item.collection_id) for item in anchors], "aligned_key_count": len(anchors), "raw_missing_count": len(anchors) - len(selected_raw), "preprocessed_missing_count": len(anchors) - len(selected_pre), "result_missing_count": len(anchors) - len(selected_results), } 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) logger.debug("Writing %d binary trace collection(s) to %s", len(history), stage_dir) 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, "capture_start_ns": int(collection.capture_start_ns), "capture_end_ns": int(collection.capture_end_ns), "trace_count": len(collection.traces), # Also in the .bin trailer; repeated here so per-combo timing is # readable without decoding the binary payload. "traces": [ { "input": int(trace.combo.input), "output": int(trace.combo.output), "capture_start_ns": int(trace.capture_start_ns), "capture_end_ns": int(trace.capture_end_ns), } for trace in 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) logger.debug("Writing %d binary result collection(s) to %s", len(history), stage_dir) 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, "processing_duration_ns": int(collection.processing_duration_ns), "collection_payload_count": len(collection.collection_payloads), "block_count": len(collection.blocks), }, indent=2, ), encoding="utf-8", ) def save_trace_history_numpy( stage_dir: Path, history: list[SweepCollection], *, index_offset: int = 0 ) -> None: """Write raw/preprocessed collections as a NumPy directory tree. ``index_offset`` continues the per-collection directory numbering across calls so a streaming recorder can append successive chunks into the same stage directory without colliding or restarting the index. """ for collection in history: validate_trace_collection(collection) stage_dir.mkdir(parents=True, exist_ok=True) logger.debug("Writing %d NumPy trace collection(s) to %s", len(history), stage_dir) for index, collection in enumerate(history, start=index_offset): 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}_o{trace.combo.output}" freq = np.asarray(trace.frequency_hz, dtype=np.float32) s11 = np.asarray(trace.s11, dtype=np.complex64) s21 = np.asarray(trace.s21, dtype=np.complex64) np.save(collection_dir / f"{tag}_freq.npy", freq) np.save(collection_dir / f"{tag}_s11.npy", s11) np.save(collection_dir / f"{tag}_s21.npy", s21) traces_meta.append( { "input": int(trace.combo.input), "output": int(trace.combo.output), "points": int(freq.size), # When each combo was measured, which in a switched matrix is # spread across the collection window rather than aligned with it. "capture_start_ns": int(trace.capture_start_ns), "capture_end_ns": int(trace.capture_end_ns), "freq_file": f"{tag}_freq.npy", "s11_file": f"{tag}_s11.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), "capture_start_ns": int(collection.capture_start_ns), "capture_end_ns": int(collection.capture_end_ns), "trace_count": len(collection.traces), "traces": traces_meta, }, indent=2, ), encoding="utf-8", ) def _save_result_payload(directory: Path, base_name: str, payload: ResultPayload) -> dict[str, Any]: """Persist every supported payload kind identically at collection and combo scope.""" meta: dict[str, Any] = {"kind": int(payload.kind), "name": payload.processing_name} def save_array(field: str, array: np.ndarray) -> None: filename = f"{base_name}_{field}.npy" np.save(directory / filename, array, allow_pickle=False) meta[f"{field}_file"] = filename if payload.kind == 1: freq = np.asarray(payload.frequency_hz, dtype=np.float32) trace = np.asarray(payload.trace, dtype=np.complex64) if freq.ndim != 1 or trace.ndim != 1 or freq.shape != trace.shape: raise ValueError("Result trace frequency and values must be equally sized 1D arrays") save_array("freq", freq) save_array("trace", trace) meta["points"] = int(freq.size) elif payload.kind == 2: save_array("scalar", np.asarray([payload.scalar_value], dtype=np.float32)) meta["scalar_value"] = float(payload.scalar_value) elif payload.kind == 3: x_axis = np.asarray(payload.image_x_axis, dtype=np.float32) y_axis = np.asarray(payload.image_y_axis, dtype=np.float32) image = np.asarray(payload.image, dtype=np.float32) if x_axis.ndim != 1 or y_axis.ndim != 1 or image.shape != (y_axis.size, x_axis.size): raise ValueError("Result image axis sizes must match its 2D matrix shape") save_array("x_axis", x_axis) save_array("y_axis", y_axis) save_array("image", image) meta.update(x_points=int(x_axis.size), y_points=int(y_axis.size)) elif payload.kind == 4: table = np.asarray(payload.table, dtype=np.float32) if table.ndim != 2: raise ValueError("Result table payload must be a 2D matrix") save_array("table", table) meta.update(rows=int(table.shape[0]), columns=int(table.shape[1])) else: raise ValueError(f"Unsupported payload kind: {payload.kind}") return meta def save_result_history_numpy( stage_dir: Path, history: list[ResultCollection], *, index_offset: int = 0 ) -> None: """Write result payloads, preserving per-stage indices across streaming chunks.""" stage_dir.mkdir(parents=True, exist_ok=True) logger.debug("Writing %d NumPy result collection(s) to %s", len(history), stage_dir) for index, collection in enumerate(history, start=index_offset): collection_dir = stage_dir / collection_dir_name(index, collection.collection_id, collection.monotonic_ns) collection_dir.mkdir(parents=True, exist_ok=False) collection_payload_meta = [] for payload_index, payload in enumerate(collection.collection_payloads): name = sanitize_path_component(payload.processing_name or "processor") base_name = f"collection_{payload_index:03d}_{name}_kind{payload.kind}" collection_payload_meta.append(_save_result_payload(collection_dir, base_name, payload)) blocks_meta = [] for block_index, block in enumerate(collection.blocks): block_dir = collection_dir / f"block_{block_index:03d}_i{block.combo.input}_o{block.combo.output}" block_dir.mkdir(parents=True, exist_ok=False) payload_meta = [] for payload_index, payload in enumerate(block.payloads): name = sanitize_path_component(payload.processing_name or "processor") base_name = f"{payload_index:03d}_{name}_kind{payload.kind}" payload_meta.append(_save_result_payload(block_dir, base_name, payload)) blocks_meta.append({ "input": int(block.combo.input), "output": int(block.combo.output), "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), "processing_duration_ns": int(collection.processing_duration_ns), "collection_payload_count": len(collection.collection_payloads), "collection_payloads": collection_payload_meta, "block_count": len(collection.blocks), "blocks": blocks_meta, }, indent=2), encoding="utf-8", )