"""Convert manual LibreVNA S11/S21 CSV captures in prog_libre to vna history JSON."""

from __future__ import annotations

import argparse
import csv
from datetime import datetime, timezone
import json
from pathlib import Path
from typing import Any

import numpy as np


def _real_imag_keys(trace_prefix: str) -> tuple[str, str]:
    return f"{trace_prefix}_Real", f"{trace_prefix}_Imaginary"


def _load_complex_trace(csv_path: Path, trace_prefix: str) -> tuple[np.ndarray, np.ndarray]:
    real_key, imag_key = _real_imag_keys(trace_prefix)
    frequencies: list[float] = []
    values: list[complex] = []

    with csv_path.open(encoding="utf-8", newline="") as handle:
        reader = csv.DictReader(handle)
        for row in reader:
            frequencies.append(float(row["Frequency"]))
            values.append(complex(float(row[real_key]), float(row[imag_key])))

    frequency_hz = np.asarray(frequencies, dtype=np.float64)
    trace = np.asarray(values, dtype=np.complex128)
    if frequency_hz.size == 0 or trace.size == 0:
        raise ValueError(f"CSV has no points: {csv_path}")
    if frequency_hz.shape != trace.shape:
        raise ValueError(f"Frequency/trace size mismatch: {csv_path}")
    return frequency_hz, trace


def _solve_one_port_osl(
    open_trace: np.ndarray,
    short_trace: np.ndarray,
    load_trace: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Solve ideal OSL one-port calibration coefficients."""
    directivity = load_trace
    open_delta = open_trace - directivity
    short_delta = short_trace - directivity
    denom = open_delta - short_delta

    source_match = np.zeros_like(directivity)
    reflection_tracking = np.ones_like(directivity)

    stable_mask = np.abs(denom) > 1e-18
    source_match[stable_mask] = (open_delta[stable_mask] + short_delta[stable_mask]) / denom[stable_mask]
    reflection_tracking[stable_mask] = open_delta[stable_mask] * (1.0 - source_match[stable_mask])
    return directivity, source_match, reflection_tracking


def _apply_one_port_osl(
    measured_trace: np.ndarray,
    directivity: np.ndarray,
    source_match: np.ndarray,
    reflection_tracking: np.ndarray,
) -> np.ndarray:
    numerator = measured_trace - directivity
    denominator = reflection_tracking + (source_match * numerator)

    corrected = np.array(numerator, copy=True)
    stable_mask = np.abs(denominator) > 1e-18
    corrected[stable_mask] = numerator[stable_mask] / denominator[stable_mask]
    return corrected


def _apply_through_calibration(measured_trace: np.ndarray, through_trace: np.ndarray) -> np.ndarray:
    corrected = np.array(measured_trace, copy=True)
    stable_mask = np.abs(through_trace) > 1e-18
    corrected[stable_mask] = measured_trace[stable_mask] / through_trace[stable_mask]
    return corrected


def _complex_to_points(values: np.ndarray) -> list[list[float]]:
    return [[float(value.real), float(value.imag)] for value in values]


def _scan_file_sort_key(csv_path: Path) -> tuple[int, str]:
    stem = csv_path.stem
    return (int(stem), stem) if stem.isdigit() else (10**9, stem)


def _load_scan_series(folder: Path, trace_prefix: str) -> tuple[np.ndarray, list[tuple[str, np.ndarray]]]:
    scan_paths = [
        path
        for path in sorted(folder.glob("*.csv"), key=_scan_file_sort_key)
        if path.stem.isdigit()
    ]
    if not scan_paths:
        raise FileNotFoundError(f"No numbered scan CSV files found in {folder}")

    base_frequency_hz: np.ndarray | None = None
    scans: list[tuple[str, np.ndarray]] = []
    for path in scan_paths:
        frequency_hz, trace = _load_complex_trace(path, trace_prefix)
        if base_frequency_hz is None:
            base_frequency_hz = frequency_hz
        elif not np.allclose(base_frequency_hz, frequency_hz, rtol=0.0, atol=1e-6):
            raise ValueError(f"Frequency axis mismatch in {path}")
        scans.append((path.name, trace))

    assert base_frequency_hz is not None
    return base_frequency_hz, scans


def _require_matching_frequency_axis(label: str, left: np.ndarray, right: np.ndarray) -> None:
    if not np.allclose(left, right, rtol=0.0, atol=1e-6):
        raise ValueError(f"{label} frequency axes do not match")


def _build_history_payload(
    *,
    source_dir: Path,
    mode: str,
    frequency_hz: np.ndarray,
    sweep_scans: list[tuple[str, np.ndarray]],
    calibrated_scans: list[tuple[str, np.ndarray]],
    reference_trace: np.ndarray,
    primary_stage: str,
    raw_record_count: int,
    preprocessed_record_count: int,
) -> dict[str, Any]:
    if len(sweep_scans) != len(calibrated_scans):
        raise ValueError("Sweep/calibrated scan counts do not match")

    sweep_history: list[dict[str, Any]] = []
    reference_points = _complex_to_points(reference_trace)
    for index, ((scan_name, sweep_trace), (cal_name, calibrated_trace)) in enumerate(
        zip(sweep_scans, calibrated_scans, strict=True)
    ):
        if scan_name != cal_name:
            raise ValueError(f"Scan ordering mismatch: {scan_name} vs {cal_name}")

        sweep_history.append(
            {
                "timestamp": float(index),
                "sweep_points": _complex_to_points(sweep_trace),
                "calibrated_points": _complex_to_points(calibrated_trace),
                "reference_points": reference_points,
                "vna_config": {
                    "mode": mode,
                    "start_freq": float(frequency_hz[0]),
                    "stop_freq": float(frequency_hz[-1]),
                    "points": int(frequency_hz.size),
                },
            }
        )

    return {
        "format": "vna-system-history-v1",
        "converter": "python_app/scripts/convert_prog_libre_manual_to_vna_history.py",
        "converted_at_utc": datetime.now(timezone.utc).isoformat(),
        "source_snapshot_dir": str(source_dir.resolve()),
        "input_index": 0,
        "output_index": 0,
        "primary_stage": primary_stage,
        "raw_record_count": int(raw_record_count),
        "preprocessed_record_count": int(preprocessed_record_count),
        "sweep_history": sweep_history,
    }


def _write_payload(output_path: Path, payload: dict[str, Any]) -> None:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")


def _rmse(left: np.ndarray, right: np.ndarray) -> float:
    return float(np.sqrt(np.mean(np.abs(left - right) ** 2)))


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Convert manual LibreVNA S11/S21 CSV captures in prog_libre to vna history JSON.",
    )
    parser.add_argument(
        "--calibration-dir",
        type=Path,
        default=Path("prog_libre/calibration"),
        help="Directory with calibration CSV files.",
    )
    parser.add_argument(
        "--raw-dir",
        type=Path,
        default=Path("prog_libre/1-6000mhz_no-calibrated_libre"),
        help="Directory with uncalibrated scan CSV files and ref.csv.",
    )
    parser.add_argument(
        "--calibrated-dir",
        type=Path,
        default=Path("prog_libre/1-6000mhz_calibrated_libre"),
        help="Directory with already calibrated scan CSV files and ref.csv.",
    )
    parser.add_argument(
        "--raw-output",
        "--s11-raw-output",
        dest="s11_raw_output",
        type=Path,
        default=Path("prog_libre/1-6000mhz_no-calibrated_libre_s11_osl_p1_vna_bscan_history.json"),
        help="Output JSON for uncalibrated S11 scans after applying OSL calibration.",
    )
    parser.add_argument(
        "--calibrated-output",
        "--s11-calibrated-output",
        dest="s11_calibrated_output",
        type=Path,
        default=Path("prog_libre/1-6000mhz_calibrated_libre_s11_passthrough_vna_bscan_history.json"),
        help="Output JSON for already calibrated S11 scans without extra calibration.",
    )
    parser.add_argument(
        "--s21-raw-output",
        dest="s21_raw_output",
        type=Path,
        default=Path("prog_libre/1-6000mhz_no-calibrated_libre_s21_through_vna_bscan_history.json"),
        help="Output JSON for uncalibrated S21 scans after applying through calibration.",
    )
    parser.add_argument(
        "--s21-calibrated-output",
        dest="s21_calibrated_output",
        type=Path,
        default=Path("prog_libre/1-6000mhz_calibrated_libre_s21_passthrough_vna_bscan_history.json"),
        help="Output JSON for already calibrated S21 scans without extra calibration.",
    )
    return parser


def main() -> None:
    args = _build_parser().parse_args()

    calibration_dir = args.calibration_dir.expanduser().resolve()
    raw_dir = args.raw_dir.expanduser().resolve()
    calibrated_dir = args.calibrated_dir.expanduser().resolve()

    s11_raw_output = args.s11_raw_output.expanduser().resolve()
    s11_calibrated_output = args.s11_calibrated_output.expanduser().resolve()
    s21_raw_output = args.s21_raw_output.expanduser().resolve()
    s21_calibrated_output = args.s21_calibrated_output.expanduser().resolve()

    s11_cal_frequency_hz, open_trace = _load_complex_trace(calibration_dir / "open_rfc18_p1.csv", "S11")
    short_frequency_hz, short_trace = _load_complex_trace(calibration_dir / "short_rfc18_p1.csv", "S11")
    load_frequency_hz, load_trace = _load_complex_trace(calibration_dir / "load_rfc18_p1.csv", "S11")
    _require_matching_frequency_axis("S11 calibration", s11_cal_frequency_hz, short_frequency_hz)
    _require_matching_frequency_axis("S11 calibration", s11_cal_frequency_hz, load_frequency_hz)

    directivity, source_match, reflection_tracking = _solve_one_port_osl(open_trace, short_trace, load_trace)

    s11_raw_frequency_hz, s11_raw_scans = _load_scan_series(raw_dir, "S11")
    s11_calibrated_frequency_hz, s11_passthrough_scans = _load_scan_series(calibrated_dir, "S11")
    _require_matching_frequency_axis("S11 raw vs calibration", s11_raw_frequency_hz, s11_cal_frequency_hz)
    _require_matching_frequency_axis("S11 calibrated vs calibration", s11_calibrated_frequency_hz, s11_cal_frequency_hz)

    s11_raw_reference_frequency_hz, s11_raw_reference = _load_complex_trace(raw_dir / "ref.csv", "S11")
    s11_calibrated_reference_frequency_hz, s11_calibrated_reference = _load_complex_trace(calibrated_dir / "ref.csv", "S11")
    _require_matching_frequency_axis("S11 raw reference vs calibration", s11_raw_reference_frequency_hz, s11_cal_frequency_hz)
    _require_matching_frequency_axis(
        "S11 calibrated reference vs calibration",
        s11_calibrated_reference_frequency_hz,
        s11_cal_frequency_hz,
    )

    s11_corrected_scans = [
        (
            scan_name,
            _apply_one_port_osl(trace, directivity, source_match, reflection_tracking),
        )
        for scan_name, trace in s11_raw_scans
    ]
    s11_corrected_reference = _apply_one_port_osl(s11_raw_reference, directivity, source_match, reflection_tracking)

    s11_raw_payload = _build_history_payload(
        source_dir=raw_dir,
        mode="s11",
        frequency_hz=s11_raw_frequency_hz,
        sweep_scans=s11_raw_scans,
        calibrated_scans=s11_corrected_scans,
        reference_trace=s11_corrected_reference,
        primary_stage="raw",
        raw_record_count=len(s11_raw_scans),
        preprocessed_record_count=len(s11_corrected_scans),
    )
    s11_calibrated_payload = _build_history_payload(
        source_dir=calibrated_dir,
        mode="s11",
        frequency_hz=s11_calibrated_frequency_hz,
        sweep_scans=s11_passthrough_scans,
        calibrated_scans=s11_passthrough_scans,
        reference_trace=s11_calibrated_reference,
        primary_stage="preprocessed",
        raw_record_count=0,
        preprocessed_record_count=len(s11_passthrough_scans),
    )

    s21_cal_frequency_hz, through_trace = _load_complex_trace(calibration_dir / "through21_rfc18.csv", "S21")

    s21_raw_frequency_hz, s21_raw_scans = _load_scan_series(raw_dir, "S21")
    s21_calibrated_frequency_hz, s21_passthrough_scans = _load_scan_series(calibrated_dir, "S21")
    _require_matching_frequency_axis("S21 raw vs calibration", s21_raw_frequency_hz, s21_cal_frequency_hz)
    _require_matching_frequency_axis("S21 calibrated vs calibration", s21_calibrated_frequency_hz, s21_cal_frequency_hz)

    s21_raw_reference_frequency_hz, s21_raw_reference = _load_complex_trace(raw_dir / "ref.csv", "S21")
    s21_calibrated_reference_frequency_hz, s21_calibrated_reference = _load_complex_trace(calibrated_dir / "ref.csv", "S21")
    _require_matching_frequency_axis("S21 raw reference vs calibration", s21_raw_reference_frequency_hz, s21_cal_frequency_hz)
    _require_matching_frequency_axis(
        "S21 calibrated reference vs calibration",
        s21_calibrated_reference_frequency_hz,
        s21_cal_frequency_hz,
    )

    s21_corrected_scans = [
        (
            scan_name,
            _apply_through_calibration(trace, through_trace),
        )
        for scan_name, trace in s21_raw_scans
    ]
    s21_corrected_reference = _apply_through_calibration(s21_raw_reference, through_trace)

    s21_raw_payload = _build_history_payload(
        source_dir=raw_dir,
        mode="s21",
        frequency_hz=s21_raw_frequency_hz,
        sweep_scans=s21_raw_scans,
        calibrated_scans=s21_corrected_scans,
        reference_trace=s21_corrected_reference,
        primary_stage="raw",
        raw_record_count=len(s21_raw_scans),
        preprocessed_record_count=len(s21_corrected_scans),
    )
    s21_calibrated_payload = _build_history_payload(
        source_dir=calibrated_dir,
        mode="s21",
        frequency_hz=s21_calibrated_frequency_hz,
        sweep_scans=s21_passthrough_scans,
        calibrated_scans=s21_passthrough_scans,
        reference_trace=s21_calibrated_reference,
        primary_stage="preprocessed",
        raw_record_count=0,
        preprocessed_record_count=len(s21_passthrough_scans),
    )

    _write_payload(s11_raw_output, s11_raw_payload)
    _write_payload(s11_calibrated_output, s11_calibrated_payload)
    _write_payload(s21_raw_output, s21_raw_payload)
    _write_payload(s21_calibrated_output, s21_calibrated_payload)

    s11_rmse_values = [
        _rmse(corrected_trace, passthrough_trace)
        for (_, corrected_trace), (_, passthrough_trace) in zip(s11_corrected_scans, s11_passthrough_scans, strict=True)
    ]
    s21_rmse_values = [
        _rmse(corrected_trace, passthrough_trace)
        for (_, corrected_trace), (_, passthrough_trace) in zip(s21_corrected_scans, s21_passthrough_scans, strict=True)
    ]
    s11_reference_rmse = _rmse(s11_corrected_reference, s11_calibrated_reference)
    s21_reference_rmse = _rmse(s21_corrected_reference, s21_calibrated_reference)

    print(
        "Converted manual prog_libre captures to vna history JSON:\n"
        f"  S11 raw output:         {s11_raw_output}\n"
        f"  S11 calibrated output:  {s11_calibrated_output}\n"
        f"  S21 raw output:         {s21_raw_output}\n"
        f"  S21 calibrated output:  {s21_calibrated_output}\n"
        f"  sweep count:            {len(s11_raw_scans)}\n"
        f"  points per sweep:       {s11_raw_frequency_hz.size}\n"
        f"  S11 calibration:        p1 ideal OSL\n"
        f"  S11 mean sweep RMSE vs provided calibrated folder: {float(np.mean(s11_rmse_values)):.6f}\n"
        f"  S11 max sweep RMSE vs provided calibrated folder:  {float(np.max(s11_rmse_values)):.6f}\n"
        f"  S11 reference RMSE vs provided calibrated ref:     {s11_reference_rmse:.6f}\n"
        f"  S21 calibration:        through21 complex division\n"
        f"  S21 mean sweep RMSE vs provided calibrated folder: {float(np.mean(s21_rmse_values)):.6f}\n"
        f"  S21 max sweep RMSE vs provided calibrated folder:  {float(np.max(s21_rmse_values)):.6f}\n"
        f"  S21 reference RMSE vs provided calibrated ref:     {s21_reference_rmse:.6f}"
    )


if __name__ == "__main__":
    main()
