#!/usr/bin/env python3
"""Fine-tune Pn/Sn picker with SNR-filtered training subsets.

The experiment compares transfer learning from the same pretrained PNSN model
under three training distributions: all CREDIT-X1 training records, records
with estimated phase SNR > 5 dB, and records with estimated phase SNR > 10 dB.
All models are evaluated on the same unfiltered test distribution.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import os
import random
import sys
from pathlib import Path
from typing import Dict, Iterable, List, Sequence

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, Loss
from scripts.reproduce_paper_stats import (
    GROUP_TO_CHANNELS,
    PHASE_TO_GROUP,
    PhasePick,
    Record,
    evaluate_outputs,
    make_specs,
    materialize_samples,
    parse_time,
    choose_phase,
    run_model,
    strip_errors,
)


COMPONENT_ORDER = ("BHE", "BHN", "BHZ")


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 record_to_dict(record: Record) -> Dict:
    return {
        "event": record.event,
        "station": record.station,
        "length": record.length,
        "delta": record.delta,
        "distance_km": record.distance_km,
        "phases": [
            {"phase": pick.phase, "index": pick.index, "source": pick.source}
            for pick in record.phases
        ],
    }


def record_from_dict(row: Dict) -> Record:
    return Record(
        event=row["event"],
        station=row["station"],
        length=int(row["length"]),
        delta=float(row["delta"]),
        distance_km=float(row["distance_km"]),
        phases=tuple(
            PhasePick(phase=p["phase"], index=int(p["index"]), source=p["source"])
            for p in row["phases"]
        ),
    )


def build_records_sequential(
    h5_path: Path,
    key_path: Path,
    split: str,
    max_events: int | None,
    cache_path: Path,
) -> List[Record]:
    if cache_path.exists():
        with cache_path.open() as f:
            cached = json.load(f)
        records = [record_from_dict(row) for row in cached["records"]]
        return records

    keys = [str(x) for x in np.load(key_path)[split]]
    if max_events is not None:
        keys = keys[:max_events]
    key_set = set(keys)
    records: List[Record] = []
    seen_events = 0
    with h5py.File(h5_path, "r") as h5:
        for i, event_key in enumerate(h5.keys()):
            if event_key not in key_set:
                continue
            seen_events += 1
            event = h5[event_key]
            for station_key in event.keys():
                station = event[station_key]
                comps = component_keys(station)
                if comps is None:
                    continue
                first = station[comps[0]]
                delta = float(first.attrs.get("delta_sec", 0.01))
                if abs(delta - 0.01) > 1e-6:
                    continue
                start_time = first.attrs.get("start_time")
                if not isinstance(start_time, str):
                    continue
                btime = parse_time(start_time)
                length = min(int(station[c].shape[0]) for c in comps)
                picks: List[PhasePick] = []
                for phase in ("Pg", "Sg", "Pn", "Sn"):
                    chosen = choose_phase(station, phase, prefer_manual=True)
                    if chosen is None:
                        continue
                    ptime, source = chosen
                    idx = int(round((parse_time(ptime) - btime).total_seconds() / delta))
                    if 0 <= idx < length:
                        picks.append(PhasePick(phase, idx, source))
                if not picks:
                    continue
                distances = []
                for phase in ("Pg", "Sg", "Pn", "Sn"):
                    for prefix in ("MANUAL.TRAVTIME", "RNN.TRAVTIME"):
                        dk = f"{prefix}.{phase}.dist_km"
                        if dk in station.attrs:
                            try:
                                distances.append(float(station.attrs[dk]))
                            except (TypeError, ValueError):
                                pass
                dist = float(np.median(distances)) if distances else float("nan")
                records.append(
                    Record(
                        event=event_key,
                        station=str(station_key),
                        length=length,
                        delta=delta,
                        distance_km=dist,
                        phases=tuple(picks),
                    )
                )
            if seen_events % 5000 == 0:
                print(
                    f"{split}: scanned {seen_events}/{len(keys)} selected events; "
                    f"records={len(records)}",
                    flush=True,
                )

    cache_path.parent.mkdir(parents=True, exist_ok=True)
    tmp = cache_path.with_suffix(".tmp")
    with tmp.open("w") as f:
        json.dump(
            {
                "split": split,
                "max_events": max_events,
                "events": len(keys),
                "records": [record_to_dict(r) for r in records],
            },
            f,
        )
    tmp.replace(cache_path)
    return records


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: Sequence[np.ndarray], phase: str, index: int) -> float | None:
    """Estimate phase SNR in dB using the local convention in utils/datapnsn.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 compute_record_snr(h5_path: Path, records: Sequence[Record], cache_path: Path) -> Dict[str, float]:
    if cache_path.exists():
        with cache_path.open() as f:
            return {k: float(v) for k, v in json.load(f).items()}

    cache_path.parent.mkdir(parents=True, exist_ok=True)
    snr_by_record: Dict[str, float] = {}
    with h5py.File(h5_path, "r") as h5:
        for i, record in enumerate(records):
            station = h5[record.event][record.station]
            comps = component_keys(station)
            if comps is None:
                continue
            waves = [station[c][:] for c in comps]
            values = []
            for pick in record.phases:
                snr = pick_snr_db(waves, pick.phase, pick.index)
                if snr is not None and math.isfinite(snr):
                    values.append(snr)
            snr_by_record[f"{record.event}/{record.station}"] = max(values) if values else -10000.0
            if i % 5000 == 0:
                print(f"computed SNR for {i}/{len(records)} records", flush=True)

    tmp = cache_path.with_suffix(".tmp")
    with tmp.open("w") as f:
        json.dump(snr_by_record, f)
    tmp.replace(cache_path)
    return snr_by_record


def filter_records_by_snr(
    records: Sequence[Record],
    snr_by_record: Dict[str, float],
    threshold: float | None,
) -> List[Record]:
    if threshold is None:
        return list(records)
    return [
        record
        for record in records
        if snr_by_record.get(f"{record.event}/{record.station}", -10000.0) > threshold
    ]


def stable_record_id(record: Record) -> str:
    return f"{record.event}/{record.station}"


def matched_records(records: Sequence[Record], n_records: int, seed: int) -> List[Record]:
    if len(records) < n_records:
        raise RuntimeError(f"Cannot draw {n_records} records from a pool of {len(records)} records.")
    ordered = sorted(records, key=stable_record_id)
    rng = np.random.default_rng(seed)
    idx = np.sort(rng.choice(len(ordered), size=n_records, replace=False))
    return [ordered[int(i)] for i in idx]


def sample_batch(
    h5_path: Path,
    records: Sequence[Record],
    seed: int,
    batch_size: int,
    length: int,
    padlen: int,
    step: int,
) -> tuple[np.ndarray, np.ndarray]:
    specs = make_specs(
        records,
        n_samples=batch_size,
        seed=seed + step * 104729,
        length=length,
        padlen=padlen,
        double_prob=0.5,
    )
    waves, _, targets, _ = materialize_samples(h5_path, records, specs, length)
    return waves, targets


def train_one(
    h5_path: Path,
    records: Sequence[Record],
    base_ckpt: Path,
    out_ckpt: Path,
    log_csv: Path,
    seed: int,
    steps: int,
    batch_size: int,
    length: int,
    padlen: int,
    lr: float,
    device: torch.device,
    resume: bool,
) -> BRNN:
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    model = BRNN().to(device)
    if resume and out_ckpt.exists():
        model.load_state_dict(torch.load(out_ckpt, map_location="cpu"))
        model.eval()
        return model

    model.load_state_dict(torch.load(base_ckpt, map_location="cpu"))
    model.train()
    loss_fn = Loss().to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
    out_ckpt.parent.mkdir(parents=True, exist_ok=True)
    log_csv.parent.mkdir(parents=True, exist_ok=True)
    with log_csv.open("w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["step", "loss"])
        for step in range(steps):
            waves, targets = sample_batch(h5_path, records, seed, batch_size, length, padlen, step)
            xb = torch.from_numpy(waves).to(device).permute(0, 2, 1)
            yb = torch.from_numpy(targets).to(device)
            out = model(xb)
            loss = loss_fn(out, yb)
            opt.zero_grad(set_to_none=True)
            loss.backward()
            opt.step()
            loss_value = float(loss.detach().cpu())
            writer.writerow([step, loss_value])
            if step % 50 == 0 or step == steps - 1:
                print(f"{out_ckpt.stem}: step {step:05d}/{steps} loss={loss_value:.3f}", flush=True)
    torch.save(model.state_dict(), out_ckpt)
    model.eval()
    return model


def metric_row(metrics: Dict, threshold: float = 0.1) -> Dict[str, float]:
    out: Dict[str, float] = {}
    for group in ("P", "S"):
        row = min(metrics["all"][group], key=lambda r: abs(r["threshold"] - threshold))
        for key in ("precision", "recall", "f1"):
            out[f"{group}_{key}"] = float(row[key])
    out["mean_f1"] = float((out["P_f1"] + out["S_f1"]) / 2.0)
    return out


def plot_summary(rows: Sequence[Dict], out: Path) -> None:
    labels = [r["label"] for r in rows]
    x = np.arange(len(labels))
    width = 0.24
    fig, ax = plt.subplots(figsize=(7.2, 4.2), dpi=220)
    for offset, key, color in [
        (-width, "P_f1", "#0072B2"),
        (0.0, "S_f1", "#D55E00"),
        (width, "mean_f1", "#009E73"),
    ]:
        ax.bar(x + offset, [r[key] for r in rows], width=width, label=key.replace("_", " "), color=color)
    ax.set_ylabel("F1 on unfiltered test set")
    ax.set_ylim(0, 1.0)
    ax.set_xticks(x)
    ax.set_xticklabels(labels)
    ax.grid(axis="y", alpha=0.25)
    ax.legend(frameon=False, ncols=3, loc="upper center", bbox_to_anchor=(0.5, 1.14))
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)


def write_metrics_table(rows: Sequence[Dict], out: Path) -> None:
    headers = [
        "Training subset",
        "Train records",
        "P precision",
        "P recall",
        "P F1",
        "S precision",
        "S recall",
        "S F1",
        "Mean F1",
    ]
    lines = ["\\begin{tabular}{lrrrrrrrr}", "\\toprule", " & ".join(headers) + " \\\\", "\\midrule"]
    for row in rows:
        lines.append(
            f"{row['label']} & {row['train_records']} & "
            f"{row['P_precision']:.3f} & {row['P_recall']:.3f} & {row['P_f1']:.3f} & "
            f"{row['S_precision']:.3f} & {row['S_recall']:.3f} & {row['S_f1']:.3f} & "
            f"{row['mean_f1']:.3f} \\\\"
        )
    lines.extend(["\\bottomrule", "\\end{tabular}", ""])
    out.write_text("\n".join(lines))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--h5", default="data/credit-x1.h5")
    parser.add_argument("--keys", default="data/creditkeys.npz")
    parser.add_argument("--base-ckpt", default="ckpt/pnsn.v3.pt")
    parser.add_argument("--out-dir", default="outputs/snr_transfer_seed20260609")
    parser.add_argument("--seed", type=int, default=20260609)
    parser.add_argument("--length", type=int, default=5120)
    parser.add_argument("--padlen", type=int, default=512)
    parser.add_argument("--train-steps", type=int, default=2000)
    parser.add_argument("--train-batch", type=int, default=16)
    parser.add_argument("--eval-samples", type=int, default=10000)
    parser.add_argument("--eval-batch", type=int, default=64)
    parser.add_argument("--max-train-events", type=int, default=0)
    parser.add_argument("--max-test-events", type=int, default=0)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--resume", action="store_true")
    parser.add_argument(
        "--no-match-train-size",
        action="store_true",
        help="Use every candidate record in each SNR pool instead of matching candidate counts.",
    )
    args = parser.parse_args()

    h5_path = Path(args.h5)
    key_path = Path(args.keys)
    base_ckpt = Path(args.base_ckpt)
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    device = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
    print(f"device={device}", flush=True)
    max_train = None if args.max_train_events == 0 else args.max_train_events
    max_test = None if args.max_test_events == 0 else args.max_test_events
    cache_tag_train = "all" if max_train is None else str(max_train)
    cache_tag_test = "all" if max_test is None else str(max_test)
    train_records = build_records_sequential(
        h5_path,
        key_path,
        "train",
        max_train,
        out_dir / f"records_train_{cache_tag_train}.json",
    )
    test_records = build_records_sequential(
        h5_path,
        key_path,
        "test",
        max_test,
        out_dir / f"records_test_{cache_tag_test}.json",
    )
    print(f"train records={len(train_records)} test records={len(test_records)}", flush=True)

    snr_by_record = compute_record_snr(h5_path, train_records, out_dir / "train_record_snr_db.json")
    snr_values = np.array(list(snr_by_record.values()), dtype=float)
    snr_summary = {
        "min": float(np.min(snr_values)),
        "median": float(np.median(snr_values)),
        "mean": float(np.mean(snr_values)),
        "p90": float(np.percentile(snr_values, 90)),
        "max": float(np.max(snr_values)),
    }

    subset_defs = [
        ("full", "Full", None),
        ("snr5", "SNR>5 dB", 5.0),
        ("snr10", "SNR>10 dB", 10.0),
    ]
    candidate_records = {
        slug: filter_records_by_snr(train_records, snr_by_record, snr_threshold)
        for slug, _, snr_threshold in subset_defs
    }
    candidate_counts = {slug: len(records) for slug, records in candidate_records.items()}
    if args.no_match_train_size:
        train_pools = candidate_records
        matched_train_records = None
    else:
        matched_train_records = min(candidate_counts.values())
        train_pools = {
            slug: matched_records(records, matched_train_records, args.seed + i * 8191)
            for i, (slug, _, _) in enumerate(subset_defs)
            for records in [candidate_records[slug]]
        }
    print(
        "candidate train records="
        + json.dumps(candidate_counts, ensure_ascii=False)
        + f"; matched_train_records={matched_train_records}",
        flush=True,
    )

    eval_specs = make_specs(
        test_records,
        n_samples=args.eval_samples,
        seed=args.seed + 17,
        length=args.length,
        padlen=args.padlen,
        double_prob=0.5,
    )
    eval_waves, eval_labels, _, eval_kinds = materialize_samples(h5_path, test_records, eval_specs, args.length)
    thresholds = [round(x, 1) for x in np.arange(0.1, 1.0, 0.1)]

    rows = []
    model_metrics = {}
    for idx, (slug, label, snr_threshold) in enumerate(subset_defs):
        subset_records = train_pools[slug]
        if not subset_records:
            raise RuntimeError(f"No records available for subset {label}")
        model = train_one(
            h5_path=h5_path,
            records=subset_records,
            base_ckpt=base_ckpt,
            out_ckpt=out_dir / f"pnsn.v3.transfer.{slug}.pt",
            log_csv=out_dir / f"transfer_loss_{slug}.csv",
            seed=args.seed + idx * 1000,
            steps=args.train_steps,
            batch_size=args.train_batch,
            length=args.length,
            padlen=args.padlen,
            lr=args.lr,
            device=device,
            resume=args.resume,
        )
        outputs = run_model(model, eval_waves, device, args.eval_batch)
        metrics = evaluate_outputs(outputs, eval_labels, eval_kinds, thresholds, min_sep=50, tolerance=100)
        model_metrics[slug] = strip_errors(metrics)
        row = {
            "slug": slug,
            "label": label if args.no_match_train_size else f"{label} matched",
            "snr_threshold_db": snr_threshold,
            "train_records": len(subset_records),
            "candidate_train_records": candidate_counts[slug],
            **metric_row(metrics),
        }
        rows.append(row)
        print(f"metrics {label}: {json.dumps(row, ensure_ascii=False)}", flush=True)

    plot_summary(rows, out_dir / "snr_transfer_f1_summary.png")
    write_metrics_table(rows, out_dir / "metrics_table.tex")
    summary = {
        "seed": args.seed,
        "device": str(device),
        "train_steps": args.train_steps,
        "train_batch": args.train_batch,
        "matched_train_records": matched_train_records,
        "candidate_train_records": candidate_counts,
        "eval_samples": args.eval_samples,
        "eval_samples_single": int(sum(k == "single" for k in eval_kinds)),
        "eval_samples_double": int(sum(k == "double" for k in eval_kinds)),
        "length_samples": args.length,
        "length_seconds": args.length * 0.01,
        "snr_definition": "max per-record phase SNR using 10*log10(std_after/std_before); P on Z, S on mean E/N; windows match utils/datapnsn.py.",
        "snr_summary_db": snr_summary,
        "rows": rows,
        "metrics_by_model": model_metrics,
        "figure": str((out_dir / "snr_transfer_f1_summary.png").resolve()),
        "table": str((out_dir / "metrics_table.tex").resolve()),
    }
    with (out_dir / "summary.json").open("w") as f:
        json.dump(summary, f, ensure_ascii=False, indent=2)
    print(json.dumps(summary, ensure_ascii=False, indent=2)[:4000], flush=True)


if __name__ == "__main__":
    main()