#!/usr/bin/env python3
"""Summarize manual P/S pick confidence and pick-level SNR."""

from __future__ import annotations

import argparse
import csv
import json
import math
import os
import sys
from collections import Counter, defaultdict
from pathlib import Path

os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("MKL_NUM_THREADS", "1")

import h5py
import matplotlib.pyplot as plt
import numpy as np
import torch


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

from models.BRNNPNSN import BRNN


DEFAULT_H5 = ROOT / "data" / "credit-x1.h5"
DEFAULT_RECORDS = [
    ROOT / "outputs" / "snr_transfer_seed20260609" / "records_train_all.json",
    ROOT / "outputs" / "snr_transfer_seed20260609" / "records_test_all.json",
]
DEFAULT_CKPT = ROOT / "outputs" / "snr_transfer_seed20260609" / "pnsn.v3.transfer.full.pt"
DEFAULT_OUT_DIR = ROOT / "outputs" / "manual_phase_confidence_snr"
PHASE_TO_GROUP = {"Pg": "P", "Pn": "P", "Sg": "S", "Sn": "S"}
GROUP_TO_CHANNELS = {"P": [1, 3], "S": [2, 4]}
COMPONENT_ORDER = ("BHE", "BHN", "BHZ")
DTYPE = np.float32


def component_keys(station) -> tuple[str, str, str] | None:
    keys = set(station.keys())
    if all(k in keys for k in COMPONENT_ORDER):
        return COMPONENT_ORDER
    by_suffix: dict[str, str] = {}
    for key in keys:
        if key.endswith("HE"):
            by_suffix["BHE"] = key
        elif key.endswith("HN"):
            by_suffix["BHN"] = key
        elif key.endswith("HZ"):
            by_suffix["BHZ"] = key
    if all(k in by_suffix for k in COMPONENT_ORDER):
        return tuple(by_suffix[k] for k in COMPONENT_ORDER)
    return None


def window_std(wave: np.ndarray, start: int, end: int) -> float | None:
    if start < 0 or end > len(wave) or end <= start:
        return None
    return float(np.std(wave[start:end]))


def pick_snr_db(waves: tuple[np.ndarray, np.ndarray, np.ndarray], phase: str, index: int) -> float | None:
    """Use the same local SNR convention as scripts/snr_transfer_experiment.py."""
    east, north, vertical = waves
    if PHASE_TO_GROUP[phase] == "P":
        pre = window_std(vertical, index - 50, index)
        aft = window_std(vertical, index, index + 50)
        if pre is None or aft is None:
            return None
        return 10.0 * math.log10((aft + 1e-6) / (pre + 1e-6))

    pre_e = window_std(east, index - 150, index)
    aft_e = window_std(east, index, index + 150)
    pre_n = window_std(north, index - 150, index)
    aft_n = window_std(north, index, index + 150)
    if pre_e is None or aft_e is None or pre_n is None or aft_n is None:
        return None
    snr_e = 10.0 * math.log10((aft_e + 1e-6) / (pre_e + 1e-6))
    snr_n = 10.0 * math.log10((aft_n + 1e-6) / (pre_n + 1e-6))
    return 0.5 * (snr_e + snr_n)


def normalize_wave(wave: np.ndarray) -> np.ndarray:
    wave = wave.astype(DTYPE, copy=False)
    wave = wave - wave.mean(axis=0, keepdims=True)
    denom = np.max(np.abs(wave), axis=0, keepdims=True) + 1e-6
    return wave / denom


def pick_window(station, comps: tuple[str, str, str], index: int, length: int) -> tuple[np.ndarray, int] | None:
    station_len = min(int(station[c].shape[0]) for c in comps)
    if station_len <= 0 or index < 0 or index >= station_len:
        return None
    start = int(np.clip(index - length // 2, 0, max(0, station_len - length)))
    stop = min(start + length, station_len)
    data = [station[c][start:stop] for c in comps]
    wave = np.stack(data, axis=1)
    if len(wave) < length:
        padded = np.zeros((length, 3), dtype=DTYPE)
        padded[: len(wave)] = wave
        wave = padded
    return normalize_wave(wave), index - start


def load_records(paths: list[Path], sample_size: int | None, seed: int) -> list[dict]:
    rows: list[dict] = []
    seen: set[tuple[str, str, str, int]] = set()
    for path in paths:
        payload = json.loads(path.read_text())
        split = payload.get("split", path.stem)
        for record in payload["records"]:
            for pick in record["phases"]:
                if str(pick.get("source", "")).startswith("MANUAL"):
                    key = (record["event"], record["station"], pick["phase"], int(pick["index"]))
                    if key in seen:
                        continue
                    seen.add(key)
                    rows.append(
                        {
                            "split": split,
                            "event": record["event"],
                            "station": record["station"],
                            "phase": pick["phase"],
                            "phase_group": PHASE_TO_GROUP[pick["phase"]],
                            "index": int(pick["index"]),
                            "distance_km": record.get("distance_km"),
                        }
                    )
    if sample_size is not None and sample_size < len(rows):
        rng = np.random.default_rng(seed)
        idx = np.sort(rng.choice(len(rows), size=sample_size, replace=False))
        rows = [rows[int(i)] for i in idx]
    return rows


def scalar_attr(value):
    if hasattr(value, "item"):
        value = value.item()
    if isinstance(value, bytes):
        return value.decode("utf-8")
    return value


def model_confidence(
    ckpt_path: Path,
    windows: np.ndarray,
    rel_indices: list[int],
    groups: list[str],
    batch_size: int,
    device: torch.device,
) -> list[float]:
    model = BRNN().to(device)
    model.load_state_dict(torch.load(ckpt_path, map_location="cpu"))
    model.eval()
    values: list[float] = []
    with torch.no_grad():
        for start in range(0, len(windows), batch_size):
            batch = torch.from_numpy(windows[start : start + batch_size]).to(device).permute(0, 2, 1)
            out = model(batch).detach().cpu().numpy()
            for j in range(out.shape[0]):
                row_i = start + j
                rel = rel_indices[row_i]
                chans = GROUP_TO_CHANNELS[groups[row_i]]
                values.append(float(out[j, chans, rel].max()))
    return values


def collect_pick_rows(
    h5_path: Path,
    rows: list[dict],
    ckpt_path: Path,
    length: int,
    batch_size: int,
    device: torch.device,
) -> list[dict]:
    out: list[dict] = []
    windows: list[np.ndarray] = []
    rel_indices: list[int] = []
    groups: list[str] = []
    missing = Counter()
    with h5py.File(h5_path, "r") as h5:
        for i, row in enumerate(rows, start=1):
            try:
                station = h5[row["event"]][row["station"]]
            except KeyError:
                missing["station"] += 1
                continue
            comps = component_keys(station)
            if comps is None:
                missing["components"] += 1
                continue
            wt_key = f"MANUAL.TRAVTIME.{row['phase']}.WT"
            if wt_key not in station.attrs:
                missing["manual_wt"] += 1
                continue
            try:
                manual_wt = float(scalar_attr(station.attrs[wt_key]))
            except (TypeError, ValueError):
                missing["manual_wt_numeric"] += 1
                continue
            waves = tuple(station[c][:] for c in comps)
            snr = pick_snr_db(waves, row["phase"], row["index"])
            if snr is None or not math.isfinite(snr):
                missing["snr"] += 1
                continue
            window = pick_window(station, comps, row["index"], length)
            if window is None:
                missing["model_window"] += 1
                continue
            wave_window, rel_index = window
            out.append(
                {
                    **row,
                    "manual_wt": manual_wt,
                    "snr_db": float(snr),
                }
            )
            windows.append(wave_window)
            rel_indices.append(rel_index)
            groups.append(row["phase_group"])
            if i % 25000 == 0:
                print(f"processed {i}/{len(rows)} manual picks; usable={len(out)}", flush=True)
    if missing:
        print("skipped:", dict(missing), flush=True)
    if out:
        conf = model_confidence(
            ckpt_path=ckpt_path,
            windows=np.stack(windows, axis=0),
            rel_indices=rel_indices,
            groups=groups,
            batch_size=batch_size,
            device=device,
        )
        for row, value in zip(out, conf):
            row["model_confidence"] = value
    return out


def describe(values: list[float]) -> dict:
    arr = np.asarray(values, dtype=float)
    if arr.size == 0:
        return {"n": 0}
    q = np.percentile(arr, [0, 5, 25, 50, 75, 95, 100])
    return {
        "n": int(arr.size),
        "mean": float(np.mean(arr)),
        "std": float(np.std(arr, ddof=1)) if arr.size > 1 else 0.0,
        "min": float(q[0]),
        "p05": float(q[1]),
        "p25": float(q[2]),
        "median": float(q[3]),
        "p75": float(q[4]),
        "p95": float(q[5]),
        "max": float(q[6]),
    }


def make_summary(rows: list[dict]) -> dict:
    summary = {
        "notes": [
            "manual_wt is MANUAL.TRAVTIME.<phase>.WT from data/README.md.",
            "model_confidence is the neural phase picker's P- or S-group probability at the manual pick sample.",
            "snr_db uses the same pick-level local window definition as scripts/snr_transfer_experiment.py.",
        ],
        "overall": {},
        "by_phase_group": {},
        "by_phase": {},
        "manual_wt_counts_by_phase_group": {},
    }
    for label, subset in [("all", rows)]:
        summary["overall"][label] = {
            "model_confidence": describe([r["model_confidence"] for r in subset]),
            "manual_wt": describe([r["manual_wt"] for r in subset]),
            "snr_db": describe([r["snr_db"] for r in subset]),
        }
    for key in ("phase_group", "phase"):
        target = summary[f"by_{key}"]
        grouped: dict[str, list[dict]] = defaultdict(list)
        for row in rows:
            grouped[row[key]].append(row)
        for name, subset in sorted(grouped.items()):
            target[name] = {
                "model_confidence": describe([r["model_confidence"] for r in subset]),
                "manual_wt": describe([r["manual_wt"] for r in subset]),
                "snr_db": describe([r["snr_db"] for r in subset]),
            }
    counts: dict[str, Counter] = defaultdict(Counter)
    for row in rows:
        counts[row["phase_group"]][str(row["manual_wt"])] += 1
    summary["manual_wt_counts_by_phase_group"] = {
        group: dict(sorted(counter.items(), key=lambda item: float(item[0])))
        for group, counter in sorted(counts.items())
    }
    return summary


def write_csv(path: Path, rows: list[dict]) -> None:
    fieldnames = [
        "split",
        "event",
        "station",
        "phase",
        "phase_group",
        "index",
        "distance_km",
        "manual_wt",
        "model_confidence",
        "snr_db",
    ]
    with path.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def scatter_panel(ax, rows: list[dict], group: str, color: str) -> None:
    subset = [r for r in rows if r["phase_group"] == group]
    x = np.asarray([r["model_confidence"] for r in subset], dtype=float)
    y = np.asarray([r["snr_db"] for r in subset], dtype=float)
    ax.scatter(
        x,
        y,
        s=9,
        c=color,
        alpha=0.28,
        linewidths=0,
        label=f"{group} picks (n={len(subset):,})",
        rasterized=True,
    )
    ax.axhline(0.0, color="#4d4d4d", lw=0.8, ls="--", alpha=0.75)
    ax.set_xlabel("Neural-network pick confidence")
    ax.set_ylabel("Pick-level SNR (dB)")
    ax.set_title(f"{group} picks")
    ax.grid(True, color="#d9d9d9", lw=0.6, alpha=0.7)
    ax.legend(frameon=False, loc="best")


def plot_scatter(path: Path, rows: list[dict]) -> None:
    colors = {"P": "#2868a8", "S": "#c44e52"}
    fig, axes = plt.subplots(2, 1, figsize=(7.2, 8.0), dpi=220, sharex=True)
    for ax, group in zip(axes, ("P", "S")):
        scatter_panel(ax, rows, group, colors[group])
    axes[-1].set_xlim(-0.02, 1.02)
    fig.suptitle("Manual Labels Scored by Neural Phase Picker: Confidence vs SNR", y=0.995)
    fig.tight_layout()
    fig.savefig(path)
    plt.close(fig)


def plot_single_group(path: Path, rows: list[dict], group: str) -> None:
    color = {"P": "#2868a8", "S": "#c44e52"}[group]
    fig, ax = plt.subplots(figsize=(7.2, 5.0), dpi=220)
    scatter_panel(ax, rows, group, color)
    ax.set_xlim(-0.02, 1.02)
    fig.tight_layout()
    fig.savefig(path)
    plt.close(fig)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--h5", type=Path, default=DEFAULT_H5)
    parser.add_argument("--records", type=Path, nargs="+", default=DEFAULT_RECORDS)
    parser.add_argument("--ckpt", type=Path, default=DEFAULT_CKPT)
    parser.add_argument("--out-dir", type=Path, default=DEFAULT_OUT_DIR)
    parser.add_argument("--sample-size", type=int, default=2000)
    parser.add_argument("--seed", type=int, default=20260626)
    parser.add_argument("--length", type=int, default=5120)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--device", default="cpu")
    args = parser.parse_args()

    args.out_dir.mkdir(parents=True, exist_ok=True)
    device = torch.device(args.device)
    candidates = load_records(args.records, args.sample_size, args.seed)
    sample_note = "all" if args.sample_size is None else f"sampled {len(candidates)}"
    print(f"loaded {sample_note} unique manual P/S picks from record caches", flush=True)
    rows = collect_pick_rows(args.h5, candidates, args.ckpt, args.length, args.batch_size, device)
    csv_path = args.out_dir / "manual_phase_confidence_snr.csv"
    summary_path = args.out_dir / "manual_phase_confidence_snr_summary.json"
    png_path = args.out_dir / "manual_phase_confidence_snr_scatter.png"
    p_png_path = args.out_dir / "manual_phase_confidence_snr_scatter_P.png"
    s_png_path = args.out_dir / "manual_phase_confidence_snr_scatter_S.png"
    write_csv(csv_path, rows)
    summary = make_summary(rows)
    summary["inputs"] = {
        "h5": str(args.h5),
        "records": [str(p) for p in args.records],
        "ckpt": str(args.ckpt),
        "sample_size": args.sample_size,
        "seed": args.seed,
        "length": args.length,
        "n_candidate_manual_picks_processed": len(candidates),
        "n_usable_manual_picks": len(rows),
    }
    summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
    plot_scatter(png_path, rows)
    plot_single_group(p_png_path, rows, "P")
    plot_single_group(s_png_path, rows, "S")
    print(f"wrote {csv_path}")
    print(f"wrote {summary_path}")
    print(f"wrote {png_path}")
    print(f"wrote {p_png_path}")
    print(f"wrote {s_png_path}")


if __name__ == "__main__":
    main()