"""Validate and visualize numpy-directory runtime snapshots."""

from __future__ import annotations

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

import numpy as np


def _load_json(path: Path) -> dict[str, Any]:
    """Load JSON object from file path."""
    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 _collection_dirs(stage_dir: Path) -> list[Path]:
    """Return sorted list of collection directories for one stage."""
    if not stage_dir.exists():
        return []
    return sorted([path for path in stage_dir.iterdir() if path.is_dir()], key=lambda path: path.name)


def _first_trace_from_raw_or_pre(collection_dir: Path) -> tuple[np.ndarray, np.ndarray, str] | None:
    """Return first trace payload from raw/preprocessed collection."""
    traces = _all_traces_from_raw_or_pre(collection_dir)
    if not traces:
        return None
    return traces[0]


def _all_traces_from_raw_or_pre(collection_dir: Path) -> list[tuple[np.ndarray, np.ndarray, str]]:
    """Load all traces from raw/preprocessed collection directory."""
    meta = _load_json(collection_dir / "meta.json")
    traces = meta.get("traces", [])
    if not isinstance(traces, list):
        raise ValueError(f"Invalid traces in {collection_dir / 'meta.json'}")

    all_traces: list[tuple[np.ndarray, np.ndarray, str]] = []
    for trace_meta in traces:
        if not isinstance(trace_meta, dict):
            raise ValueError(f"Invalid trace record in {collection_dir / 'meta.json'}")

        freq_file = str(trace_meta.get("freq_file", ""))
        s11_file = str(trace_meta.get("s11_file", ""))
        s21_file = str(trace_meta.get("s21_file", ""))
        freq = np.load(collection_dir / freq_file)
        s11 = np.load(collection_dir / s11_file)
        s21 = np.load(collection_dir / s21_file)
        if freq.shape != s11.shape:
            raise ValueError(f"Shape mismatch freq/s11 in {collection_dir}")
        if freq.shape != s21.shape:
            raise ValueError(f"Shape mismatch freq/s21 in {collection_dir}")
        if not (
            np.isfinite(freq).all()
            and np.isfinite(np.real(s11)).all()
            and np.isfinite(np.imag(s11)).all()
            and np.isfinite(np.real(s21)).all()
            and np.isfinite(np.imag(s21)).all()
        ):
            raise ValueError(f"Non-finite values in {collection_dir}")

        label = f"i{int(trace_meta.get('input', 0))}_o{int(trace_meta.get('output', 0))}"
        all_traces.append((np.asarray(freq, dtype=np.float64), np.asarray(s21, dtype=np.complex128), label))

    return all_traces


def _first_trace_from_results(collection_dir: Path) -> tuple[np.ndarray, np.ndarray, str] | None:
    """Return first trace payload from results collection."""
    traces = _all_traces_from_results(collection_dir)
    if not traces:
        return None
    return traces[0]


def _all_traces_from_results(collection_dir: Path) -> list[tuple[np.ndarray, np.ndarray, str]]:
    """Load all trace-like payloads from results collection directory."""
    meta = _load_json(collection_dir / "meta.json")
    blocks = meta.get("blocks", [])
    if not isinstance(blocks, list):
        raise ValueError(f"Invalid blocks in {collection_dir / 'meta.json'}")

    all_traces: list[tuple[np.ndarray, np.ndarray, str]] = []
    for block in blocks:
        if not isinstance(block, dict):
            continue
        block_dir_name = str(block.get("dir", ""))
        block_dir = collection_dir / block_dir_name
        payloads = block.get("payloads", [])
        if not isinstance(payloads, list):
            continue
        for payload in payloads:
            if not isinstance(payload, dict):
                continue
            kind = int(payload.get("kind", 0))
            if kind != 1:
                continue
            freq_file = str(payload.get("freq_file", ""))
            trace_file = str(payload.get("trace_file", ""))
            freq = np.load(block_dir / freq_file)
            trace = np.load(block_dir / trace_file)
            if freq.shape != trace.shape:
                raise ValueError(f"Shape mismatch freq/trace in {block_dir}")
            if not (
                np.isfinite(freq).all()
                and np.isfinite(np.real(trace)).all()
                and np.isfinite(np.imag(trace)).all()
            ):
                raise ValueError(f"Non-finite values in {block_dir}")

            label = (
                f"i{int(block.get('input', 0))}_o{int(block.get('output', 0))}_"
                f"{str(payload.get('name', 'processor'))}"
            )
            all_traces.append((np.asarray(freq, dtype=np.float64), np.asarray(trace, dtype=np.complex128), label))

    return all_traces


def _validate_stage(stage_dir: Path, stage: str) -> tuple[list[int], list[tuple[np.ndarray, np.ndarray, str]]]:
    """Validate one stage directory and collect representative traces."""
    collection_ids: list[int] = []
    traces: list[tuple[np.ndarray, np.ndarray, str]] = []
    for collection_dir in _collection_dirs(stage_dir):
        meta = _load_json(collection_dir / "meta.json")
        collection_ids.append(int(meta.get("collection_id", -1)))
        if stage in {"raw", "preprocessed"}:
            trace = _first_trace_from_raw_or_pre(collection_dir)
        else:
            trace = _first_trace_from_results(collection_dir)
        if trace is not None:
            traces.append(trace)
    return collection_ids, traces


def _compare_two(
    name: str,
    first: tuple[np.ndarray, np.ndarray, str],
    second: tuple[np.ndarray, np.ndarray, str],
) -> None:
    """Print numerical difference metrics for two traces."""
    freq_a, data_a, label_a = first
    freq_b, data_b, label_b = second
    same_shape = freq_a.shape == freq_b.shape == data_a.shape == data_b.shape
    if not same_shape:
        print(f"[{name}] different shapes: {freq_a.shape}/{freq_b.shape} {data_a.shape}/{data_b.shape}")
        return

    are_equal = np.array_equal(data_a, data_b)
    diff = data_a - data_b
    max_abs_diff = float(np.max(np.abs(diff)))
    l2_diff = float(np.linalg.norm(diff))
    print(
        f"[{name}] compare first two traces: {label_a} vs {label_b}, "
        f"equal={are_equal}, max_abs_diff={max_abs_diff:.6g}, l2_diff={l2_diff:.6g}"
    )


def _plot_two(
    stage: str,
    first: tuple[np.ndarray, np.ndarray, str],
    second: tuple[np.ndarray, np.ndarray, str],
    output_dir: Path,
) -> None:
    """Plot magnitude comparison for two traces."""
    try:
        import matplotlib.pyplot as plt
    except Exception as exc:  # noqa: BLE001
        print(f"[{stage}] matplotlib is not available, plot skipped: {exc}")
        return

    output_dir.mkdir(parents=True, exist_ok=True)
    freq_a, data_a, label_a = first
    freq_b, data_b, label_b = second
    if freq_a.shape != freq_b.shape or data_a.shape != data_b.shape:
        print(f"[{stage}] shapes differ, plot skipped")
        return

    y_a = 20.0 * np.log10(np.maximum(np.abs(data_a), 1e-12))
    y_b = 20.0 * np.log10(np.maximum(np.abs(data_b), 1e-12))

    fig, ax = plt.subplots(figsize=(10, 4))
    ax.plot(freq_a, y_a, linewidth=1.4, label=f"collection#1 {label_a}")
    ax.plot(freq_b, y_b, linewidth=1.4, label=f"collection#2 {label_b}")
    ax.set_title(f"{stage}: first two collections")
    ax.set_xlabel("X axis")
    ax.set_ylabel("Magnitude dB")
    ax.grid(True, alpha=0.3)
    ax.legend()
    fig.tight_layout()

    png_path = output_dir / f"{stage}_first_two.png"
    fig.savefig(png_path, dpi=150)
    plt.close(fig)
    print(f"[{stage}] plot saved: {png_path}")


def _plot_all_states_for_one_collection(stage: str, collection_dir: Path, output_dir: Path) -> None:
    """Plot all switch-state traces available in one collection."""
    try:
        import matplotlib.pyplot as plt
    except Exception as exc:  # noqa: BLE001
        print(f"[{stage}] matplotlib is not available, all-states plot skipped: {exc}")
        return

    if stage in {"raw", "preprocessed"}:
        traces = _all_traces_from_raw_or_pre(collection_dir)
    else:
        traces = _all_traces_from_results(collection_dir)

    if not traces:
        print(f"[{stage}] no trace payloads found in {collection_dir.name}, all-states plot skipped")
        return

    output_dir.mkdir(parents=True, exist_ok=True)
    meta = _load_json(collection_dir / "meta.json")
    collection_id = int(meta.get("collection_id", -1))

    fig, ax = plt.subplots(figsize=(11, 5))
    for freq, data, label in traces:
        y = 20.0 * np.log10(np.maximum(np.abs(data), 1e-12))
        ax.plot(freq, y, linewidth=1.2, label=label)
    ax.set_title(f"{stage}: all switch states in one collection (id={collection_id})")
    ax.set_xlabel("X axis")
    ax.set_ylabel("Magnitude dB")
    ax.grid(True, alpha=0.3)
    ax.legend(fontsize=8, ncol=2)
    fig.tight_layout()

    png_path = output_dir / f"{stage}_all_states_one_collection.png"
    fig.savefig(png_path, dpi=150)
    plt.close(fig)
    print(f"[{stage}] all-states plot saved: {png_path} (traces={len(traces)})")


def _run(snapshot_dir: Path, output_dir: Path) -> None:
    """Execute snapshot validation and plotting workflow."""
    manifest_path = snapshot_dir / "manifest.json"
    if manifest_path.exists():
        manifest = _load_json(manifest_path)
        print(
            f"Snapshot: {snapshot_dir}\n"
            f"selection_mode={manifest.get('selection_mode')} "
            f"raw={manifest.get('raw_collections')} "
            f"pre={manifest.get('preprocessed_collections')} "
            f"res={manifest.get('result_collections')}"
        )
    else:
        print(f"Snapshot: {snapshot_dir} (manifest.json is missing)")

    stages = ("raw", "preprocessed", "results")
    for stage in stages:
        stage_dir = snapshot_dir / stage
        collection_dirs = _collection_dirs(stage_dir)
        collection_ids, traces = _validate_stage(stage_dir, stage)
        duplicate_id_count = len(collection_ids) - len(set(collection_ids))
        print(
            f"[{stage}] collections={len(collection_ids)}, "
            f"duplicate_ids={duplicate_id_count}, "
            f"trace_samples={len(traces)}"
        )
        if len(traces) >= 2:
            _compare_two(stage, traces[0], traces[1])
            _plot_two(stage, traces[0], traces[1], output_dir)
        else:
            print(f"[{stage}] not enough trace-like collections to compare/plot (need >= 2)")

        if collection_dirs:
            _plot_all_states_for_one_collection(stage, collection_dirs[-1], output_dir)
        else:
            print(f"[{stage}] no collections for all-states plot")


def main() -> None:
    """CLI entrypoint."""
    parser = argparse.ArgumentParser(description="Validate and visualize runtime numpy snapshot collections.")
    parser.add_argument(
        "snapshot_dir",
        type=Path,
        help="Path to snapshot directory (contains raw/preprocessed/results).",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=None,
        help="Directory for output plots (default: <snapshot_dir>/inspection_plots).",
    )
    args = parser.parse_args()

    snapshot_dir = args.snapshot_dir.expanduser().resolve()
    if not snapshot_dir.exists():
        raise FileNotFoundError(f"Snapshot directory not found: {snapshot_dir}")

    output_dir = (
        args.output_dir.expanduser().resolve()
        if args.output_dir is not None
        else snapshot_dir / "inspection_plots"
    )
    _run(snapshot_dir, output_dir)


if __name__ == "__main__":
    main()