radar_system/python_app/storage/npz/store.py

"""Concrete :class:`StoreApi` implementation backed by NPZ files."""

from __future__ import annotations

from contextlib import suppress
from datetime import datetime, timezone
import json
import logging
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

logger = logging.getLogger(__name__)


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:
    """Append aligned raw/preprocessed/result collections to a snapshot dir in chunks.

    A streaming alternative to :meth:`NpzStore.save_runtime_snapshot_numpy` for long
    recordings: the caller flushes small chunks as measurements arrive and frees them,
    so memory stays flat instead of buffering every measurement. The on-disk layout is
    identical (``raw``/``preprocessed``/``results`` subtrees), and per-collection
    directory indices continue across chunks via a running offset per stage.
    """

    __slots__ = ("_dir", "_raw_written", "_preprocessed_written", "_result_written")

    def __init__(self, snapshot_dir: Path) -> None:
        self._dir = snapshot_dir
        self._raw_written = 0
        self._preprocessed_written = 0
        self._result_written = 0

    @property
    def directory(self) -> Path:
        return self._dir

    @property
    def result_count(self) -> int:
        """Number of result collections written to disk so far."""
        return self._result_written

    def append(
        self,
        raw: list[SweepCollection],
        preprocessed: list[SweepCollection],
        results: list[ResultCollection],
    ) -> None:
        """Write one chunk of (already aligned) collections, continuing each stage's index."""
        save_trace_history_numpy(self._dir / "raw", raw, index_offset=self._raw_written)
        save_trace_history_numpy(
            self._dir / "preprocessed", preprocessed, index_offset=self._preprocessed_written
        )
        save_result_history_numpy(self._dir / "results", results, index_offset=self._result_written)
        self._raw_written += len(raw)
        self._preprocessed_written += len(preprocessed)
        self._result_written += len(results)


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)
        logger.debug("NpzStore rooted at %s", self._root_dir)

    @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,
                    "capture_start_ns": int(trace.capture_start_ns),
                    "capture_end_ns": int(trace.capture_end_ns),
                    "freq_key": freq_key,
                    "s11_key": s11_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,
        }

        # Write both files to temporary paths first, then atomically rename so a
        # crash never leaves an .npz without its meta (or vice versa).
        npz_tmp = npz_path.with_name(npz_path.name + ".tmp")
        meta_tmp = meta_path.with_name(meta_path.name + ".tmp")
        try:
            with npz_tmp.open("wb") as npz_file:
                np.savez(npz_file, **payload)
            meta_tmp.write_text(json.dumps(meta, indent=2), encoding="utf-8")
            npz_tmp.replace(npz_path)
            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)
            for tmp_path in (npz_tmp, meta_tmp):
                with suppress(OSError):
                    tmp_path.unlink(missing_ok=True)
            raise
        logger.info(
            "Saved set %s/%s/%s (%d traces) to %s", kind, radar_key, set_name, len(combo_records), npz_path
        )

    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():
            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:
            meta = json.loads(meta_path.read_text(encoding="utf-8"))
            arrays = np.load(npz_path, allow_pickle=False)

            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,
                        capture_start_ns=int(combo.get("capture_start_ns", 0)),
                        capture_end_ns=int(combo.get("capture_end_ns", 0)),
                    )
                )

            collection = 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)),
            )
            logger.debug("Loaded set %s (%d traces)", set_label, len(traces))
            return collection
        except (json.JSONDecodeError, 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)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_bytes(serialize_trace_collection(collection, RAW_MAGIC))
        logger.info("Exported set %s/%s/%s bundle to %s", kind, radar_key, set_name, output_path)
        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.now(timezone.utc).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:])
        logger.info("Saved binary runtime snapshot (last_n=%d) to %s", last_n, snapshot_dir)
        return snapshot_dir

    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],
        last_n: int,
        *,
        name_prefix: str = "",
    ) -> 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 = _compose_snapshot_stem(snapshot_name, name_prefix)

        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)
        logger.info(
            "Saved NumPy runtime snapshot to %s (raw=%d preprocessed=%d results=%d, anchor=%s)",
            snapshot_dir,
            len(selected_raw),
            len(selected_preprocessed),
            len(selected_results),
            selection_summary.get("anchor_stage"),
        )
        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.now(timezone.utc).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)
        logger.info(
            "Saved VNA history JSON to %s (input=%d output=%d channel=%s sweeps=%d)",
            output_path,
            int(input_index),
            int(output_index),
            channel,
            summary["sweep_count"],
        )
        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",
        name_prefix: str = "",
    ) -> 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 = _compose_snapshot_stem(output_name, name_prefix)
        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:
            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}_{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]
        logger.info(
            "Saved %d VNA history JSON file(s) to %s (channel=%s)", len(output_paths), output_dir, channel
        )
        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

    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


__all__ = ["NpzStore", "radar_key_from_config"]