"""Concrete :class:`StoreApi` implementation backed by NPZ files.""" from __future__ import annotations from datetime import datetime import json from pathlib import Path from typing import Any import numpy as np from python_app.models.dataset_model import ComboKey, ResultCollection, SweepCollection, TraceData from python_app.storage.npz.paths import radar_key_from_config, sanitize_path_component from python_app.storage.npz.serialize import PREPROC_MAGIC, RAW_MAGIC, serialize_trace_collection from python_app.storage.npz.snapshot_numpy import ( save_result_history_binary, save_result_history_numpy, save_trace_history_binary, save_trace_history_numpy, select_aligned_histories, ) from python_app.storage.npz.vna_history_json import build_vna_history_payload from python_app.storage.store_api import StoreApi class NpzStore(StoreApi): """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 self._root_dir.mkdir(parents=True, exist_ok=True) @staticmethod def _vna_json_output_dir(output_root_dir: Path, output_stem: str) -> Path: """Return per-export directory for VNA history JSON files.""" return output_root_dir / output_stem def save_set(self, kind: str, radar_key: str, set_name: str, collection: SweepCollection) -> None: """Persist named preprocess set as NPZ and metadata JSON.""" set_dir = self._set_dir(kind, radar_key) set_dir.mkdir(parents=True, exist_ok=True) npz_path = set_dir / f"{set_name}.npz" meta_path = set_dir / f"{set_name}.json" payload: dict[str, np.ndarray] = {} 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}" s11_key = f"s11_{suffix}" s21_key = f"s21_{suffix}" payload[freq_key] = np.asarray(trace.frequency_hz, dtype=np.float32) payload[s11_key] = np.asarray(trace.s11, dtype=np.complex64) payload[s21_key] = np.asarray(trace.s21, dtype=np.complex64) combo_records.append( { "input": trace.combo.input, "output": trace.combo.output, "freq_key": freq_key, "s11_key": s11_key, "s21_key": s21_key, } ) np.savez(npz_path, **payload) 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, } meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8") def load_set(self, kind: str, radar_key: str, set_name: str) -> SweepCollection: """Load named preprocess set from NPZ representation.""" 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() or not meta_path.exists(): raise FileNotFoundError(f"Missing set files for {kind}/{radar_key}/{set_name}") meta = json.loads(meta_path.read_text(encoding="utf-8")) arrays = np.load(npz_path) traces: list[TraceData] = [] for combo in meta["combos"]: freq = np.asarray(arrays[combo["freq_key"]], dtype=np.float32) s11 = np.asarray(arrays[combo["s11_key"]], dtype=np.complex64) s21 = np.asarray(arrays[combo["s21_key"]], dtype=np.complex64) traces.append( TraceData( combo=ComboKey(input=int(combo["input"]), output=int(combo["output"])), frequency_hz=freq, s11=s11, s21=s21, ) ) return SweepCollection( collection_id=int(meta["collection_id"]), monotonic_ns=int(meta["monotonic_ns"]), traces=traces, capture_start_ns=int(meta.get("capture_start_ns", 0)), capture_end_ns=int(meta.get("capture_end_ns", 0)), ) 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) output_path.parent.mkdir(parents=True, exist_ok=True) output_path.write_bytes(serialize_trace_collection(collection, RAW_MAGIC)) return output_path def save_runtime_snapshot( self, output_dir: Path, raw_history: list[SweepCollection], preprocessed_history: list[SweepCollection], result_history: list[ResultCollection], last_n: int, ) -> Path: """Save historical runtime collections using binary on-disk format.""" if last_n <= 0: raise ValueError("last_n must be > 0") output_dir.mkdir(parents=True, exist_ok=True) timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S") snapshot_dir = output_dir / f"snapshot_{timestamp}" snapshot_dir.mkdir(parents=True, exist_ok=True) save_trace_history_binary(snapshot_dir / "raw", raw_history[-last_n:], RAW_MAGIC) save_trace_history_binary(snapshot_dir / "preprocessed", preprocessed_history[-last_n:], PREPROC_MAGIC) save_result_history_binary(snapshot_dir / "results", result_history[-last_n:]) return 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], last_n: int, ) -> tuple[Path, dict[str, Any]]: """Save historical runtime collections in NumPy tree format.""" if last_n <= 0: raise ValueError("last_n must be > 0") snapshot_stem = snapshot_name.strip() or datetime.utcnow().strftime("snapshot_%Y%m%d_%H%M%S") snapshot_stem = sanitize_path_component(snapshot_stem) output_root_dir.mkdir(parents=True, exist_ok=True) snapshot_dir = output_root_dir / snapshot_stem if snapshot_dir.exists(): raise FileExistsError(f"Snapshot directory already exists: {snapshot_dir}") snapshot_dir.mkdir(parents=True, exist_ok=False) selected_raw, selected_preprocessed, selected_results, selection_summary = select_aligned_histories( raw_history, preprocessed_history, result_history, last_n, ) save_trace_history_numpy(snapshot_dir / "raw", selected_raw) save_trace_history_numpy(snapshot_dir / "preprocessed", selected_preprocessed) save_result_history_numpy(snapshot_dir / "results", selected_results) if selection_summary.get("anchor_stage") == "results": expected = min(int(last_n), len(result_history)) if len(selected_results) != expected: raise RuntimeError( "Snapshot results selection invariant failed: " f"expected={expected}, selected={len(selected_results)}" ) selection_summary["raw_count"] = len(selected_raw) selection_summary["preprocessed_count"] = len(selected_preprocessed) selection_summary["result_count"] = len(selected_results) selection_summary["snapshot_stem"] = snapshot_stem selection_summary["snapshot_dir"] = str(snapshot_dir) return snapshot_dir, selection_summary def save_runtime_vna_history_json( self, output_root_dir: Path, output_name: str, raw_history: list[SweepCollection], preprocessed_history: list[SweepCollection], result_history: list[ResultCollection], last_n: int, *, input_index: int = 0, output_index: int = 0, channel: str = "s21", primary_stage: str = "preprocessed", ) -> tuple[Path, dict[str, Any]]: """Save runtime history as vna_system-compatible JSON file.""" if last_n <= 0: raise ValueError("last_n must be > 0") output_stem = output_name.strip() or datetime.utcnow().strftime("snapshot_%Y%m%d_%H%M%S") output_stem = sanitize_path_component(output_stem) output_root_dir.mkdir(parents=True, exist_ok=True) output_dir = self._vna_json_output_dir(output_root_dir, output_stem) output_dir.mkdir(parents=True, exist_ok=True) output_path = output_dir / f"{output_stem}_vna_bscan_history.json" if output_path.exists(): raise FileExistsError(f"Output JSON file already exists: {output_path}") selected_raw, selected_preprocessed, selected_results, selection_summary = select_aligned_histories( raw_history, preprocessed_history, result_history, last_n, ) payload = build_vna_history_payload( selected_raw, selected_preprocessed, selected_results, input_index=input_index, output_index=output_index, channel=channel, primary_stage=primary_stage, ) output_path.write_text( json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8", ) summary: dict[str, Any] = dict(selection_summary) summary["raw_count"] = len(selected_raw) summary["preprocessed_count"] = len(selected_preprocessed) summary["result_count"] = len(selected_results) summary["raw_record_count"] = int(payload.get("raw_record_count", 0)) summary["preprocessed_record_count"] = int(payload.get("preprocessed_record_count", 0)) summary["sweep_count"] = len(payload.get("sweep_history", [])) summary["input_index"] = int(input_index) summary["output_index"] = int(output_index) summary["channel"] = str(channel) summary["primary_stage"] = str(primary_stage) summary["output_stem"] = output_stem summary["output_dir"] = str(output_dir) summary["output_path"] = str(output_path) return output_path, summary 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], last_n: int, *, channel: str = "s21", primary_stage: str = "preprocessed", ) -> tuple[list[Path], dict[str, Any]]: """Save one runtime VNA-history JSON per available combo.""" if last_n <= 0: raise ValueError("last_n must be > 0") output_stem = output_name.strip() or datetime.utcnow().strftime("snapshot_%Y%m%d_%H%M%S") output_stem = sanitize_path_component(output_stem) output_root_dir.mkdir(parents=True, exist_ok=True) output_dir = self._vna_json_output_dir(output_root_dir, output_stem) output_dir.mkdir(parents=True, exist_ok=True) selected_raw, selected_preprocessed, selected_results, selection_summary = select_aligned_histories( raw_history, preprocessed_history, result_history, last_n, ) combos = sorted( { (int(trace.combo.input), int(trace.combo.output)) for collection in [*selected_raw, *selected_preprocessed] for trace in collection.traces } ) if not combos: 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}_{channel}_vna_bscan_history.json" ) if output_path.exists(): raise FileExistsError(f"Output JSON file already exists: {output_path}") payloads.append( build_vna_history_payload( selected_raw, selected_preprocessed, selected_results, input_index=input_index, output_index=output_index, channel=channel, primary_stage=primary_stage, ) ) output_paths.append(output_path) for output_path, payload in zip(output_paths, payloads, strict=True): output_path.write_text( json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8", ) summary: dict[str, Any] = dict(selection_summary) summary["raw_count"] = len(selected_raw) summary["preprocessed_count"] = len(selected_preprocessed) summary["result_count"] = len(selected_results) summary["preprocessed_record_count"] = int(sum(int(payload.get("preprocessed_record_count", 0)) for payload in payloads)) summary["combo_count"] = len(combos) summary["combos"] = [list(combo) for combo in combos] summary["channel"] = str(channel) summary["primary_stage"] = str(primary_stage) summary["output_stem"] = output_stem summary["output_dir"] = str(output_dir) summary["output_paths"] = [str(path) for path in output_paths] return output_paths, summary def _set_dir(self, kind: str, radar_key: str) -> Path: """Return directory for set kind and radar key.""" return self._root_dir / kind / radar_key __all__ = ["NpzStore", "radar_key_from_config"]