"""Preprocess sets and datasets backed by NPZ files.""" from __future__ import annotations from contextlib import suppress from datetime import datetime, timezone import json import logging import os from pathlib import Path import shutil from typing import Any from uuid import uuid4 from zipfile import BadZipFile import numpy as np from python_app.file_io import make_directory, rename_no_replace, sync_directory, temporary_file from python_app.models.dataset_model import ComboKey, ResultCollection, SweepCollection, TraceData from python_app.file_io import sanitize_path_component from python_app.orchestration.shm.decoder import RAW_MAGIC from python_app.storage.npz.serialize import serialize_trace_collection, validate_trace_collection from python_app.storage.npz.snapshot_numpy import ( save_result_history_numpy, save_trace_history_numpy, save_processing_settings, atomic_snapshot_directory, ) from python_app.storage.npz.vna_history_json import build_vna_history_payload logger = logging.getLogger(__name__) _SET_METADATA_KEY = "metadata_json" def _compose_snapshot_stem(name: str, name_prefix: str) -> str: """Compose a sanitized snapshot stem from an optional prefix and the operator's name. The name falls back to a UTC timestamp when blank, so a save is always uniquely named. """ base = name.strip() or datetime.now(timezone.utc).strftime("snapshot_%Y%m%d_%H%M%S") return sanitize_path_component(f"{name_prefix}_{base}" if name_prefix else base) class SnapshotStreamWriter: """Publish complete collections incrementally, retaining them if a later write fails.""" def __init__(self, snapshot_dir: Path) -> None: self._dir = snapshot_dir self._counts = {"raw": 0, "preprocessed": 0, "results": 0} @property def directory(self) -> Path: return self._dir @property def result_count(self) -> int: return self._counts["results"] def append( self, raw: list[SweepCollection], preprocessed: list[SweepCollection], results: list[ResultCollection], *, processing_settings: dict[int, dict] | None = None, ) -> None: """Append a chunk; after failure, only the unwritten collections may be retried.""" save_processing_settings(self._dir, processing_settings) for stage, history in (("raw", raw), ("preprocessed", preprocessed)): stage_dir = self._dir / stage stage_dir.mkdir(exist_ok=True) for collection in history: save_trace_history_numpy(stage_dir, [collection], index_offset=self._counts[stage]) self._counts[stage] += 1 stage_dir = self._dir / "results" stage_dir.mkdir(exist_ok=True) for result in results: save_result_history_numpy(stage_dir, [result], index_offset=self.result_count) self._counts["results"] += 1 class NpzStore: """Persist preprocess sets and runtime snapshots using NumPy files.""" def __init__(self, root_dir: Path) -> None: """Create store rooted at `root_dir`.""" self._root_dir = root_dir make_directory(self._root_dir, durable=True) logger.debug("NpzStore rooted at %s", self._root_dir) def save_set( self, kind: str, radar_key: str, set_name: str, collection: SweepCollection, *, metadata: dict[str, Any] | None = None, ) -> None: """Persist a named set with optional JSON metadata embedded in the NPZ.""" self._validate_set_name(set_name) validate_trace_collection(collection, require_finite=True) if metadata is not None and not isinstance(metadata, dict): raise ValueError("Set user metadata must be a JSON object") set_dir = self._set_dir(kind, radar_key) make_directory(set_dir, durable=True) npz_path = set_dir / f"{set_name}.npz" meta_path = set_dir / f"{set_name}.json" payload: dict[str, Any] = {} combo_records: list[dict[str, str | int]] = [] for trace in collection.traces: suffix = f"i{trace.combo.input}_o{trace.combo.output}" freq_key = f"freq_{suffix}" s21_key = f"s21_{suffix}" payload[freq_key] = np.asarray(trace.frequency_hz, dtype=np.float32) payload[s21_key] = np.asarray(trace.s21, dtype=np.complex64) combo_records.append( { "input": trace.combo.input, "output": trace.combo.output, "capture_start_ns": int(trace.capture_start_ns), "capture_end_ns": int(trace.capture_end_ns), "freq_key": freq_key, "s21_key": s21_key, } ) meta = { "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), "combos": combo_records, "metadata": {} if metadata is None else metadata, } # Metadata in the NPZ makes a single atomic replacement authoritative. The # JSON remains a readable sidecar and legacy NPZ/JSON pairs remain loadable. metadata_json = json.dumps(meta, indent=2, allow_nan=False) payload[_SET_METADATA_KEY] = np.asarray(metadata_json) npz_tmp = meta_tmp = backup_path = None archive_replaced = False try: npz_tmp = temporary_file(npz_path) meta_tmp = temporary_file(meta_path) backup_path = temporary_file(npz_path) if npz_path.exists() else None if backup_path is not None: backup_path.unlink() os.link(npz_path, backup_path) with npz_tmp.open("wb") as npz_file: np.savez(npz_file, **payload) npz_file.flush() os.fsync(npz_file.fileno()) with meta_tmp.open("w", encoding="utf-8") as meta_file: meta_file.write(metadata_json) meta_file.flush() os.fsync(meta_file.fileno()) npz_tmp.replace(npz_path) archive_replaced = True meta_tmp.replace(meta_path) except BaseException: logger.exception("Failed to save set %s/%s/%s; rolling back temp files", kind, radar_key, set_name) if archive_replaced: if backup_path is None: npz_path.unlink(missing_ok=True) else: try: backup_path.replace(npz_path) except OSError: logger.exception("Set rollback failed; previous archive retained at %s", backup_path) backup_path = None # retain the recovery copy rather than deleting it below raise sync_directory(set_dir) raise finally: for tmp_path in (npz_tmp, meta_tmp, backup_path): if tmp_path is None: continue with suppress(OSError): tmp_path.unlink(missing_ok=True) sync_directory(set_dir) logger.info( "Saved set %s/%s/%s (%d traces) to %s", kind, radar_key, set_name, len(combo_records), npz_path ) def save_sets(self, radar_key: str, set_name: str, collections: dict[str, SweepCollection]) -> None: """Save related sets together, restoring their previous files if any save fails.""" self._validate_set_name(set_name) for collection in collections.values(): validate_trace_collection(collection, require_finite=True) backup_dir = self._root_dir / f".sets-{uuid4().hex}" backup_dir.mkdir() backups: dict[Path, Path | None] = {} retain_backups = False try: for kind in collections: set_dir = self._set_dir(kind, radar_key) for extension in ("npz", "json"): destination = set_dir / f"{set_name}.{extension}" backup = backup_dir / f"{kind}.{extension}" if destination.exists() else None if backup is not None: os.link(destination, backup) backups[destination] = backup try: for kind, collection in collections.items(): self.save_set(kind, radar_key, set_name, collection) except BaseException as save_error: retain_backups = True # preserve recovery files if rollback itself is interrupted rollback_failed = False for destination, backup in backups.items(): try: if backup is None: destination.unlink(missing_ok=True) else: backup.replace(destination) except OSError: rollback_failed = True logger.exception("Failed to restore previous set file %s", destination) for directory in {path.parent for path in backups if path.parent.exists()}: try: sync_directory(directory) except OSError: rollback_failed = True logger.exception("Failed to sync restored set directory %s", directory) retain_backups = rollback_failed if retain_backups: raise RuntimeError( f"Saving set group '{set_name}' failed and rollback was incomplete; " f"previous files are retained at {backup_dir}" ) from save_error raise finally: if not retain_backups: with suppress(OSError): shutil.rmtree(backup_dir) def load_set(self, kind: str, radar_key: str, set_name: str) -> SweepCollection: """Load named preprocess set from NPZ representation.""" return self.load_set_with_metadata(kind, radar_key, set_name)[0] def load_set_with_metadata( self, kind: str, radar_key: str, set_name: str, ) -> tuple[SweepCollection, dict[str, Any]]: """Load a named set and its user metadata; older sets have empty metadata.""" self._validate_set_name(set_name) set_dir = self._set_dir(kind, radar_key) npz_path = set_dir / f"{set_name}.npz" meta_path = set_dir / f"{set_name}.json" if not npz_path.exists(): logger.error("Missing set files for %s/%s/%s", kind, radar_key, set_name) raise FileNotFoundError(f"Missing set files for {kind}/{radar_key}/{set_name}") set_label = f"{kind}/{radar_key}/{set_name}" try: with np.load(npz_path, allow_pickle=False) as arrays: meta = json.loads( str(arrays[_SET_METADATA_KEY].item()) if _SET_METADATA_KEY in arrays else meta_path.read_text(encoding="utf-8") ) if not isinstance(meta, dict): raise ValueError("Set metadata must be a JSON object") metadata = meta.get("metadata", {}) if not isinstance(metadata, dict): raise ValueError("Set user metadata must be a JSON object") json.dumps(metadata, allow_nan=False) traces: list[TraceData] = [] for combo in meta["combos"]: traces.append( TraceData( combo=ComboKey(input=combo["input"], output=combo["output"]), frequency_hz=np.asarray(arrays[combo["freq_key"]], dtype=np.float32), s21=np.asarray(arrays[combo["s21_key"]], dtype=np.complex64), capture_start_ns=combo.get("capture_start_ns", 0), capture_end_ns=combo.get("capture_end_ns", 0), ) ) collection = SweepCollection( collection_id=meta["collection_id"], monotonic_ns=meta["monotonic_ns"], traces=traces, capture_start_ns=meta.get("capture_start_ns", 0), capture_end_ns=meta.get("capture_end_ns", 0), ) validate_trace_collection(collection, require_finite=True) logger.debug("Loaded set %s (%d traces)", set_label, len(traces)) return collection, metadata except (OSError, BadZipFile, EOFError, KeyError, TypeError, ValueError) as exc: logger.exception("Corrupted preprocess set %s", set_label) raise RuntimeError(f"Corrupted preprocess set {set_label}: {exc}") from exc def list_sets(self, kind: str, radar_key: str) -> list[str]: """List available set names for `(kind, radar_key)`.""" set_dir = self._set_dir(kind, radar_key) if not set_dir.exists(): return [] return sorted(path.stem for path in set_dir.glob("*.npz")) def has_combo_coverage(self, kind: str, radar_key: str, set_name: str, combos: list[ComboKey]) -> bool: """Validate that named set covers all required switch combinations.""" collection = self.load_set(kind, radar_key, set_name) existing = {(trace.combo.input, trace.combo.output) for trace in collection.traces} required = {(combo.input, combo.output) for combo in combos} return required.issubset(existing) def export_set_bundle(self, kind: str, radar_key: str, set_name: str, output_path: Path) -> Path: """Export named set as binary collection bundle for C++ preprocessing stage.""" collection = self.load_set(kind, radar_key, set_name) return self._export_collection_bundle(collection, output_path) def _export_collection_bundle(self, collection: SweepCollection, output_path: Path) -> Path: """Publish a complete binary trace bundle using one atomic replacement.""" validate_trace_collection(collection, require_finite=True) output_path.parent.mkdir(parents=True, exist_ok=True) temporary = temporary_file(output_path) try: temporary.write_bytes(serialize_trace_collection(collection, RAW_MAGIC)) temporary.replace(output_path) finally: temporary.unlink(missing_ok=True) logger.info("Exported %d traces to bundle %s", len(collection.traces), output_path) return output_path def snapshot_directory( self, output_root_dir: Path, snapshot_name: str, *, name_prefix: str = "" ) -> Path: """Return the directory a snapshot save would create for these inputs. Lets callers check ``.exists()`` *before* acquiring data (e.g. the disk recorder arms a run only once the destination is free), so a name clash is reported up front instead of after the measurements are collected. Note a blank ``snapshot_name`` resolves to a fresh timestamp each call, so this is meaningful only for explicit names — exactly the case that can collide. """ return output_root_dir / _compose_snapshot_stem(snapshot_name, name_prefix) def create_snapshot_stream( self, output_root_dir: Path, snapshot_name: str, *, name_prefix: str = "" ) -> SnapshotStreamWriter: """Create an empty snapshot directory and return a chunked stream writer for it. Raises ``FileExistsError`` if the directory already exists (same guard as :meth:`save_runtime_snapshot_numpy`), so a name clash is reported at arm time. """ snapshot_dir = self.snapshot_directory(output_root_dir, snapshot_name, name_prefix=name_prefix) output_root_dir.mkdir(parents=True, exist_ok=True) if snapshot_dir.exists(): raise FileExistsError(f"Snapshot directory already exists: {snapshot_dir}") snapshot_dir.mkdir(parents=True, exist_ok=False) logger.info("Opened streaming snapshot directory %s", snapshot_dir) return SnapshotStreamWriter(snapshot_dir) def save_runtime_snapshot_numpy( self, output_root_dir: Path, snapshot_name: str, raw_history: list[SweepCollection], preprocessed_history: list[SweepCollection], result_history: list[ResultCollection], *, name_prefix: str = "", processing_settings: dict[int, dict] | None = None, ) -> Path: """Save exactly the supplied collections and their processing revisions.""" snapshot_dir = self.snapshot_directory(output_root_dir, snapshot_name, name_prefix=name_prefix) with atomic_snapshot_directory(snapshot_dir) as staging: save_processing_settings(staging, processing_settings) save_trace_history_numpy(staging / "raw", raw_history) save_trace_history_numpy(staging / "preprocessed", preprocessed_history) save_result_history_numpy(staging / "results", result_history) logger.info("Saved NumPy runtime snapshot to %s", snapshot_dir) return snapshot_dir def save_runtime_vna_history_json_batch( self, output_root_dir: Path, output_name: str, raw_history: list[SweepCollection], preprocessed_history: list[SweepCollection], result_history: list[ResultCollection], *, primary_stage: str = "preprocessed", name_prefix: str = "", ) -> list[Path]: """Save one runtime VNA-history JSON per available combo.""" output_stem = _compose_snapshot_stem(output_name, name_prefix) output_root_dir.mkdir(parents=True, exist_ok=True) output_dir = output_root_dir / output_stem combos = sorted( { (int(trace.combo.input), int(trace.combo.output)) for collection in [*raw_history, *preprocessed_history] for trace in collection.traces } ) if not combos: logger.warning("No raw/preprocessed combos found in runtime history for VNA JSON batch export") raise ValueError("No matching raw/preprocessed traces were found in runtime history for any combo.") output_paths: list[Path] = [] payloads: list[dict[str, Any]] = [] for input_index, output_index in combos: output_path = output_dir / ( f"{output_stem}_i{input_index}_o{output_index}_s21_vna_bscan_history.json" ) if output_path.exists(): raise FileExistsError(f"Output JSON file already exists: {output_path}") payloads.append( build_vna_history_payload( raw_history, preprocessed_history, result_history, input_index=input_index, output_index=output_index, primary_stage=primary_stage, ) ) output_paths.append(output_path) output_existed = output_dir.exists() output_dir.mkdir(exist_ok=True) published: list[Path] = [] try: # Publish complete files without overwriting a concurrent export. for output_path, payload in zip(output_paths, payloads, strict=True): temporary = temporary_file(output_path) try: with temporary.open("w", encoding="utf-8") as handle: json.dump(payload, handle, ensure_ascii=False, indent=2, allow_nan=False) handle.flush() os.fsync(handle.fileno()) rename_no_replace(temporary, output_path) published.append(output_path) finally: temporary.unlink(missing_ok=True) sync_directory(output_dir) except BaseException: for path in published: path.unlink() if not output_existed: with suppress(OSError): output_dir.rmdir() raise logger.info( "Saved %d S21 VNA history JSON file(s) to %s", len(output_paths), output_dir ) return output_paths def _set_dir(self, kind: str, radar_key: str) -> Path: """Return directory for set kind and radar key.""" self._validate_set_name(kind) self._validate_set_name(radar_key) return self._root_dir / kind / radar_key @staticmethod def _validate_set_name(value: str) -> None: """Keep caller-provided set components inside their storage directory.""" if not value or value in {".", ".."} or Path(value).name != value or "\\" in value: raise ValueError(f"Invalid set path component: {value!r}") def preview_png_dir(self, kind: str, set_name: str, radar_key: str) -> Path: """Return directory for preview PNGs of one set/radar variant. Lives under ``preview_png/`` inside the store so saved graphs sit next to the data they describe, grouped by set name then radar variant. The set name is operator-supplied, so it is sanitized; ``radar_key`` is already filesystem-safe by construction. """ return self._root_dir / "preview_png" / kind / sanitize_path_component(set_name) / radar_key