#!/usr/bin/env python3
"""Compare M>=3-only and all-magnitude transfer training at matched manual-pick counts."""

from __future__ import annotations

import argparse
import csv
import json
import math
import os
import random
import sys
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import 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 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 (
    DTYPE,
    GROUP_TO_CHANNELS,
    PHASE_TO_CHANNEL,
    PHASE_TO_GROUP,
    find_peaks,
    labels_to_target,
    normalize_wave,
)

COMPONENTS = ("BHE", "BHN", "BHZ")
PHASES = ("Pg", "Sg", "Pn", "Sn")


@dataclass(frozen=True)
class ManualPickSample:
    split: str
    event: str
    station: str
    phase: str
    phase_group: str
    index: int
    length: int
    magnitude: float
    distance_km: float | None


def component_keys(station) -> tuple[str, str, str] | None:
    keys = set(station.keys())
    if all(k in keys for k in COMPONENTS):
        return COMPONENTS
    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 COMPONENTS):
        return tuple(by_suffix[k] for k in COMPONENTS)
    return None


def load_record_cache(path: Path) -> list[dict]:
    return json.loads(path.read_text())["records"]


def event_magnitudes(h5_path: Path, events: Sequence[str]) -> dict[str, float]:
    out: dict[str, float] = {}
    with h5py.File(h5_path, "r") as h5:
        for event in sorted(set(events)):
            out[event] = float(h5[event].attrs.get("event_magnitude", float("nan")))
    return out


def manual_pick_samples(
    h5_path: Path,
    record_cache: Path,
    split: str,
    cache_path: Path,
) -> list[ManualPickSample]:
    if cache_path.exists():
        payload = json.loads(cache_path.read_text())
        return [ManualPickSample(**row) for row in payload["samples"]]

    records = load_record_cache(record_cache)
    mags = event_magnitudes(h5_path, [r["event"] for r in records])
    samples: list[ManualPickSample] = []
    for record in records:
        mag = mags[record["event"]]
        for pick in record["phases"]:
            if pick.get("source") != "MANUAL":
                continue
            phase = pick["phase"]
            samples.append(
                ManualPickSample(
                    split=split,
                    event=record["event"],
                    station=record["station"],
                    phase=phase,
                    phase_group=PHASE_TO_GROUP[phase],
                    index=int(pick["index"]),
                    length=int(record["length"]),
                    magnitude=mag,
                    distance_km=record.get("distance_km"),
                )
            )
    cache_path.parent.mkdir(parents=True, exist_ok=True)
    cache_path.write_text(
        json.dumps({"split": split, "samples": [asdict(s) for s in samples]}, indent=2),
        encoding="utf-8",
    )
    return samples


def stable_pick_id(sample: ManualPickSample) -> str:
    return f"{sample.event}/{sample.station}/{sample.phase}/{sample.index:08d}"


def matched_pick_subset(samples: Sequence[ManualPickSample], n: int, seed: int) -> list[ManualPickSample]:
    if len(samples) < n:
        raise RuntimeError(f"Cannot sample {n} picks from pool of {len(samples)}.")
    ordered = sorted(samples, key=stable_pick_id)
    rng = np.random.default_rng(seed)
    idx = np.sort(rng.choice(len(ordered), size=n, replace=False))
    return [ordered[int(i)] for i in idx]


def crop_start(sample: ManualPickSample, length: int, padlen: int, rng: np.random.Generator | None) -> int | None:
    if sample.length < length:
        return None
    if rng is None:
        offset = length // 2
    else:
        offset = int(rng.integers(padlen, max(padlen + 1, length - padlen)))
    start = sample.index - offset
    return int(np.clip(start, 0, sample.length - length))


def load_pick_window(
    h5,
    sample: ManualPickSample,
    length: int,
    padlen: int,
    rng: np.random.Generator | None,
) -> tuple[np.ndarray, np.ndarray, int] | None:
    start = crop_start(sample, length, padlen, rng)
    if start is None:
        return None
    station = h5[sample.event][sample.station]
    comps = component_keys(station)
    if comps is None:
        return None
    wave = np.stack([station[c][start : start + length] for c in comps], axis=1)
    if wave.shape[0] != length:
        return None
    wave = normalize_wave(wave)
    rel = sample.index - start
    if not 0 <= rel < length:
        return None
    target = labels_to_target([(sample.phase, sample.phase_group, rel)], length)
    return wave, target, rel


def sample_batch(
    h5_path: Path,
    samples: Sequence[ManualPickSample],
    seed: int,
    batch_size: int,
    length: int,
    padlen: int,
    step: int,
) -> tuple[np.ndarray, np.ndarray]:
    rng = np.random.default_rng(seed + step * 104729)
    waves: list[np.ndarray] = []
    targets: list[np.ndarray] = []
    valid = [s for s in samples if s.length >= length]
    with h5py.File(h5_path, "r") as h5:
        while len(waves) < batch_size:
            sample = valid[int(rng.integers(0, len(valid)))]
            loaded = load_pick_window(h5, sample, length, padlen, rng)
            if loaded is None:
                continue
            wave, target, _ = loaded
            waves.append(wave)
            targets.append(target)
    return np.stack(waves, axis=0), np.stack(targets, axis=0)


def train_one(
    h5_path: Path,
    samples: Sequence[ManualPickSample],
    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,
    init_mode: str = "transfer",
) -> 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

    if init_mode == "transfer":
        model.load_state_dict(torch.load(base_ckpt, map_location="cpu"))
    elif init_mode != "scratch":
        raise ValueError(f"Unknown init_mode={init_mode!r}")
    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)
    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, samples, 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 run_model(model: BRNN, waves: np.ndarray, device: torch.device, batch_size: int) -> np.ndarray:
    model.eval()
    outputs: list[np.ndarray] = []
    with torch.no_grad():
        for start in range(0, len(waves), batch_size):
            xb = torch.from_numpy(waves[start : start + batch_size]).to(device).permute(0, 2, 1)
            outputs.append(model(xb).detach().cpu().numpy())
    return np.concatenate(outputs, axis=0)


def materialize_eval(
    h5_path: Path,
    samples: Sequence[ManualPickSample],
    length: int,
    padlen: int,
) -> tuple[np.ndarray, list[int], list[str], list[str]]:
    waves: list[np.ndarray] = []
    rel_indices: list[int] = []
    groups: list[str] = []
    phases: list[str] = []
    with h5py.File(h5_path, "r") as h5:
        for i, sample in enumerate(samples):
            loaded = load_pick_window(h5, sample, length, padlen, rng=None)
            if loaded is None:
                continue
            wave, _target, rel = loaded
            waves.append(wave)
            rel_indices.append(rel)
            groups.append(sample.phase_group)
            phases.append(sample.phase)
            if i % 2000 == 0:
                print(f"materialized eval {i:,}/{len(samples):,}", flush=True)
    return np.stack(waves, axis=0), rel_indices, groups, phases


def evaluate_pick_recall(
    outputs: np.ndarray,
    rel_indices: Sequence[int],
    groups: Sequence[str],
    phases: Sequence[str],
    thresholds: Sequence[float],
    tolerance: int,
    min_sep: int,
) -> dict:
    out = {"thresholds": list(thresholds), "by_group": {}, "by_phase": {}}
    for scope_name, labels in [("by_group", groups), ("by_phase", phases)]:
        for label in sorted(set(labels)):
            rows = []
            idxs = [i for i, v in enumerate(labels) if v == label]
            if scope_name == "by_group":
                channel_groups = {label: GROUP_TO_CHANNELS[label]}
            else:
                channel_groups = {label: [PHASE_TO_CHANNEL[label]]}
            chans = channel_groups[label]
            for threshold in thresholds:
                tp = fp = fn = 0
                for i in idxs:
                    prob = outputs[i, chans, :].max(axis=0)
                    pred = [pidx for pidx, _score in find_peaks(prob, threshold, min_sep)]
                    matched = any(abs(pidx - rel_indices[i]) <= tolerance for pidx in pred)
                    if matched:
                        tp += 1
                        fp += max(0, len(pred) - 1)
                    else:
                        fn += 1
                        fp += len(pred)
                precision = tp / (tp + fp) if tp + fp else 0.0
                recall = tp / (tp + fn) if tp + fn else 0.0
                f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0
                rows.append(
                    {
                        "threshold": float(threshold),
                        "n": len(idxs),
                        "tp": int(tp),
                        "fp": int(fp),
                        "fn": int(fn),
                        "precision": float(precision),
                        "recall": float(recall),
                        "f1": float(f1),
                    }
                )
            out[scope_name][label] = rows
    for threshold in thresholds:
        rows = []
        tp = fp = fn = 0
        for group, group_rows in out["by_group"].items():
            row = next(r for r in group_rows if abs(r["threshold"] - threshold) < 1e-9)
            tp += row["tp"]
            fp += row["fp"]
            fn += row["fn"]
        precision = tp / (tp + fp) if tp + fp else 0.0
        recall = tp / (tp + fn) if tp + fn else 0.0
        f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0
        rows.append(
            {
                "threshold": float(threshold),
                "n": int(tp + fn),
                "tp": int(tp),
                "fp": int(fp),
                "fn": int(fn),
                "precision": float(precision),
                "recall": float(recall),
                "f1": float(f1),
            }
        )
        out.setdefault("combined", []).extend(rows)
    return out


def summarize_counts(samples: Sequence[ManualPickSample]) -> dict:
    by_phase: dict[str, int] = {}
    by_group: dict[str, int] = {}
    events = set()
    for sample in samples:
        by_phase[sample.phase] = by_phase.get(sample.phase, 0) + 1
        by_group[sample.phase_group] = by_group.get(sample.phase_group, 0) + 1
        events.add(sample.event)
    return {
        "picks": len(samples),
        "events": len(events),
        "by_phase": dict(sorted(by_phase.items())),
        "by_group": dict(sorted(by_group.items())),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--h5", type=Path, default=ROOT / "data" / "credit-x1.h5")
    parser.add_argument("--base-ckpt", type=Path, default=ROOT / "ckpt" / "pnsn.v3.pt")
    parser.add_argument("--train-record-cache", type=Path, default=ROOT / "outputs" / "snr_transfer_seed20260609" / "records_train_all.json")
    parser.add_argument("--test-record-cache", type=Path, default=ROOT / "outputs" / "snr_transfer_seed20260609" / "records_test_all.json")
    parser.add_argument("--out-dir", type=Path, default=ROOT / "outputs" / "magnitude_pick_transfer_seed20260628")
    parser.add_argument("--seed", type=int, default=20260628)
    parser.add_argument("--mag-threshold", type=float, default=3.0)
    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-batch", type=int, default=64)
    parser.add_argument("--max-eval-picks", type=int, default=0)
    parser.add_argument("--eval-all-magnitudes", action="store_true")
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--thresholds", type=float, nargs="+", default=[0.1, 0.2, 0.3, 0.5])
    parser.add_argument("--tolerance-samples", type=int, default=100)
    parser.add_argument("--min-sep", type=int, default=50)
    parser.add_argument("--resume", action="store_true")
    parser.add_argument("--device", default=None)
    args = parser.parse_args()

    if args.device is None:
        device = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
    else:
        device = torch.device(args.device)
    print(f"device={device}", flush=True)
    args.out_dir.mkdir(parents=True, exist_ok=True)

    train_samples = manual_pick_samples(args.h5, args.train_record_cache, "train", args.out_dir / "manual_train_picks.json")
    test_samples = manual_pick_samples(args.h5, args.test_record_cache, "test", args.out_dir / "manual_test_picks.json")
    train_m3 = [s for s in train_samples if s.magnitude >= args.mag_threshold and s.length >= args.length]
    train_all_valid = [s for s in train_samples if s.length >= args.length]
    if args.eval_all_magnitudes:
        eval_samples = [s for s in test_samples if s.length >= args.length]
        eval_tag = "all_magnitudes"
    else:
        eval_samples = [s for s in test_samples if s.magnitude >= args.mag_threshold and s.length >= args.length]
        eval_tag = "mag_ge_3"
    if args.max_eval_picks > 0:
        eval_samples = matched_pick_subset(eval_samples, min(args.max_eval_picks, len(eval_samples)), args.seed + 17)
    train_budget = len(train_m3)
    train_all_matched = matched_pick_subset(train_all_valid, train_budget, args.seed)

    subsets = {
        "all_mag_matched": train_all_matched,
        "mag_ge_3": sorted(train_m3, key=stable_pick_id),
    }
    print(f"train budget manual picks={train_budget}; eval {eval_tag} manual picks={len(eval_samples)}", flush=True)
    models: dict[str, BRNN] = {}
    for slug, subset in subsets.items():
        print(f"training {slug}: {summarize_counts(subset)}", flush=True)
        models[slug] = train_one(
            args.h5,
            subset,
            args.base_ckpt,
            args.out_dir / f"pnsn.v3.transfer.{slug}.pt",
            args.out_dir / f"train_log_{slug}.csv",
            args.seed,
            args.train_steps,
            args.train_batch,
            args.length,
            args.padlen,
            args.lr,
            device,
            args.resume,
        )

    waves, rel_indices, groups, phases = materialize_eval(args.h5, sorted(eval_samples, key=stable_pick_id), args.length, args.padlen)
    rows = []
    metrics_by_model = {}
    for slug, model in models.items():
        outputs = run_model(model, waves, device, args.eval_batch)
        metrics = evaluate_pick_recall(outputs, rel_indices, groups, phases, args.thresholds, args.tolerance_samples, args.min_sep)
        metrics_by_model[slug] = metrics
        for scope, labels in [("by_group", metrics["by_group"]), ("by_phase", metrics["by_phase"])]:
            for label, vals in labels.items():
                for row in vals:
                    rows.append({"model": slug, "scope": scope, "label": label, **row})
        for row in metrics["combined"]:
            rows.append({"model": slug, "scope": "combined", "label": "all", **row})

    suffix = "" if eval_tag == "mag_ge_3" else f"_eval_{eval_tag}"
    csv_path = args.out_dir / f"magnitude_pick_recall_metrics{suffix}.csv"
    with csv_path.open("w", newline="") as f:
        fieldnames = ["model", "scope", "label", "threshold", "n", "tp", "fp", "fn", "precision", "recall", "f1"]
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    summary = {
        "experiment": "manual_pick_magnitude_transfer",
        "notes": [
            "Training samples are individual MANUAL.TRAVTIME Pg/Sg/Pn/Sn picks.",
            "The all-magnitude and M>=3 training conditions are matched to the same manual-pick count.",
            "Evaluation uses only M>=3 test-set manual picks, one centered window per pick.",
        ],
        "config": {
            "seed": args.seed,
            "mag_threshold": args.mag_threshold,
            "train_steps": args.train_steps,
            "train_batch": args.train_batch,
            "max_eval_picks": args.max_eval_picks,
            "eval_all_magnitudes": args.eval_all_magnitudes,
            "length": args.length,
            "padlen": args.padlen,
            "thresholds": args.thresholds,
            "tolerance_samples": args.tolerance_samples,
            "min_sep": args.min_sep,
        },
        "counts": {
            "train_all_manual": summarize_counts(train_samples),
            "train_magnitude_ge_3_valid": summarize_counts(train_m3),
            "train_all_matched": summarize_counts(train_all_matched),
            "eval_valid": summarize_counts(eval_samples),
        },
        "metrics": metrics_by_model,
    }
    summary_path = args.out_dir / f"summary{suffix}.json"
    summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print(f"wrote {summary_path}")
    print(f"wrote {csv_path}")


if __name__ == "__main__":
    main()