radar_system/python_app/scripts/convert_legacy_preprocess_sets.py

#!/usr/bin/env python3
"""Convert legacy preprocess set storage into the current two-channel format."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import sys
from typing import Any

import numpy as np

PROJECT_ROOT = Path(__file__).resolve().parents[2]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from python_app.models.dataset_model import ComboKey, SweepCollection, TraceData
from python_app.storage.npz_store import NpzStore


LEGACY_KIND_MAP: dict[str, str] = {
    "calibration": "s21_calibration",
    "reference": "s21_reference",
    "s21_calibration": "s21_calibration",
    "s21_reference": "s21_reference",
    "s11_open": "s11_open",
    "s11_short": "s11_short",
    "s11_load": "s11_load",
    "s11_reference": "s11_reference",
}

S21_ONLY_TARGET_KINDS = {"s21_calibration", "s21_reference"}


def _build_parser() -> argparse.ArgumentParser:
    """Return the argument parser for the legacy preprocess-set conversion CLI."""
    parser = argparse.ArgumentParser(
        description=(
            "Convert old preprocess-set storage from a legacy python_app/data tree into the "
            "current format required by the new GUI/runtime."
        )
    )
    parser.add_argument(
        "source_data_dir",
        type=Path,
        help="Path to legacy python_app/data directory from the old project copy.",
    )
    parser.add_argument(
        "output_data_dir",
        type=Path,
        help="Destination directory where converted sets will be written in the new format.",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Allow overwriting already converted destination sets.",
    )
    return parser


def _load_json(path: Path) -> dict[str, Any]:
    """Read a JSON file and return its top-level object, rejecting non-object roots."""
    payload = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise ValueError(f"JSON root must be object: {path}")
    return payload


def _read_combo_position(combo_payload: dict[str, Any], *, primary_key: str, alias_key: str) -> int:
    """Return a switch position from a combo record, accepting the primary or alias key."""
    if primary_key in combo_payload:
        return int(combo_payload[primary_key])
    if alias_key in combo_payload:
        return int(combo_payload[alias_key])
    raise KeyError(f"Missing combo position field: {primary_key}/{alias_key}")


def _combo_suffix(input_pos: int, output_pos: int) -> str:
    """Return the per-combo array-name suffix (e.g. ``i0_o1``) used in legacy NPZ keys."""
    return f"i{input_pos}_o{output_pos}"


def _load_array(arrays: Any, key: str, *, dtype: np.dtype[Any], label: str) -> np.ndarray:
    """Return a flattened array of the given dtype from an NPZ mapping, by key."""
    if key not in arrays:
        raise KeyError(f"Missing {label} array '{key}' in NPZ archive")
    return np.asarray(arrays[key], dtype=dtype).reshape(-1)


def _load_legacy_collection(meta_path: Path, npz_path: Path, *, target_kind: str) -> SweepCollection:
    """Build a SweepCollection from a legacy meta/NPZ pair for the given target kind.

    Reads per-combo frequency, S21 and (when present) S11 arrays. For S21-only target
    kinds a missing S11 is filled with zeros; for any other kind a missing S11 is an error.
    """
    meta = _load_json(meta_path)
    combos_payload = meta.get("combos")
    if not isinstance(combos_payload, list):
        raise ValueError(f"Expected 'combos' list in {meta_path}")

    with np.load(npz_path) as arrays:
        traces: list[TraceData] = []
        for combo_payload in combos_payload:
            if not isinstance(combo_payload, dict):
                raise ValueError(f"Expected combo object in {meta_path}")

            input_pos = _read_combo_position(combo_payload, primary_key="input", alias_key="input_pos")
            output_pos = _read_combo_position(combo_payload, primary_key="output", alias_key="output_pos")
            suffix = _combo_suffix(input_pos, output_pos)

            freq_key = str(combo_payload.get("freq_key") or f"freq_{suffix}")
            s21_key = str(combo_payload.get("s21_key") or f"s21_{suffix}")
            s11_key = str(combo_payload.get("s11_key") or f"s11_{suffix}")

            frequency_hz = _load_array(arrays, freq_key, dtype=np.float32, label="frequency")
            s21 = _load_array(arrays, s21_key, dtype=np.complex64, label="S21")

            if frequency_hz.shape != s21.shape:
                raise ValueError(f"Frequency/S21 shape mismatch in {npz_path}: {frequency_hz.shape} vs {s21.shape}")

            if s11_key in arrays:
                s11 = _load_array(arrays, s11_key, dtype=np.complex64, label="S11")
            elif target_kind in S21_ONLY_TARGET_KINDS:
                s11 = np.zeros_like(s21, dtype=np.complex64)
            else:
                raise KeyError(
                    f"Missing S11 array '{s11_key}' in {npz_path}; "
                    f"cannot convert target kind '{target_kind}' without real S11 data"
                )

            if frequency_hz.shape != s11.shape:
                raise ValueError(f"Frequency/S11 shape mismatch in {npz_path}: {frequency_hz.shape} vs {s11.shape}")

            traces.append(
                TraceData(
                    combo=ComboKey(input=input_pos, output=output_pos),
                    frequency_hz=frequency_hz,
                    s11=s11,
                    s21=s21,
                )
            )

    return SweepCollection(
        collection_id=int(meta.get("collection_id", 0)),
        monotonic_ns=int(meta.get("monotonic_ns", 0)),
        traces=traces,
    )


def _convert_one_set(
    store: NpzStore,
    *,
    output_root: Path,
    source_kind: str,
    target_kind: str,
    radar_key: str,
    meta_path: Path,
    overwrite: bool,
) -> None:
    """Convert a single legacy set (meta + NPZ) and save it under the target kind.

    Raises if the companion NPZ is missing, or if the destination already exists and
    ``overwrite`` is False.
    """
    set_name = meta_path.stem
    npz_path = meta_path.with_suffix(".npz")
    if not npz_path.exists():
        raise FileNotFoundError(f"Missing NPZ archive for set '{set_name}': {npz_path}")

    target_dir = output_root / target_kind / radar_key
    target_json = target_dir / f"{set_name}.json"
    target_npz = target_dir / f"{set_name}.npz"
    if not overwrite and (target_json.exists() or target_npz.exists()):
        raise FileExistsError(f"Destination set already exists: {target_dir / set_name}")

    collection = _load_legacy_collection(meta_path, npz_path, target_kind=target_kind)
    store.save_set(target_kind, radar_key, set_name, collection)
    print(
        f"[converted] {source_kind}/{radar_key}/{set_name} -> "
        f"{target_kind}/{radar_key}/{set_name} (traces={len(collection.traces)})"
    )


def main() -> int:
    """Walk the legacy data tree, convert every recognized set, and print a summary.

    Returns 0 on full success, 1 if any set failed to convert, or 2 if no convertible
    sets were found.
    """
    parser = _build_parser()
    args = parser.parse_args()

    source_root = args.source_data_dir.expanduser().resolve()
    output_root = args.output_data_dir.expanduser().resolve()

    if not source_root.exists():
        raise FileNotFoundError(f"Source data directory does not exist: {source_root}")
    if source_root == output_root:
        raise ValueError("Source and output directories must be different")

    store = NpzStore(output_root)
    converted_count = 0
    skipped_kind_count = 0
    error_messages: list[str] = []

    for source_kind_dir in sorted(path for path in source_root.iterdir() if path.is_dir()):
        source_kind = source_kind_dir.name
        target_kind = LEGACY_KIND_MAP.get(source_kind)
        if target_kind is None:
            skipped_kind_count += 1
            print(f"[skip-kind] {source_kind_dir}")
            continue

        for radar_key_dir in sorted(path for path in source_kind_dir.iterdir() if path.is_dir()):
            radar_key = radar_key_dir.name
            for meta_path in sorted(radar_key_dir.glob("*.json")):
                try:
                    _convert_one_set(
                        store,
                        output_root=output_root,
                        source_kind=source_kind,
                        target_kind=target_kind,
                        radar_key=radar_key,
                        meta_path=meta_path,
                        overwrite=bool(args.overwrite),
                    )
                    converted_count += 1
                except Exception as exc:  # noqa: BLE001
                    error_messages.append(f"{meta_path}: {exc}")
                    print(f"[error] {meta_path}: {exc}")

    print(
        "\nConversion summary:\n"
        f"  source root:     {source_root}\n"
        f"  output root:     {output_root}\n"
        f"  converted sets:  {converted_count}\n"
        f"  skipped kinds:   {skipped_kind_count}\n"
        f"  errors:          {len(error_messages)}"
    )

    if error_messages:
        return 1
    if converted_count == 0:
        print("No convertible preprocess sets were found.")
        return 2
    return 0


if __name__ == "__main__":
    raise SystemExit(main())