#!/usr/bin/env python3
"""Magnitude-filtered transfer experiment for manual phase-pick recall.

Compares a model fine-tuned on all-magnitude manual phase examples against a
model fine-tuned on M>=threshold manual phase examples. Training pools are
matched by manual phase-pick count, not by event count or station-record count.
Evaluation uses only manual picks from M>=threshold test events.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import os
import random
import sys
from dataclasses import 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 find_peaks


PHASES = ("Pg", "Sg", "Pn", "Sn")
PHASE_TO_CHANNEL = {"Pg": 1, "Sg": 2, "Pn": 3, "Sn": 4}
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


@dataclass(frozen=True)
class ManualPick:
    phase: str
    index: int


@dataclass(frozen=True)
class RecordInfo:
    event: str
    station: str
    length: int
    magnitude: float
    picks: tuple[ManualPick, ...]


@dataclass(frozen=True)
class PhaseExample:
    record_id: str
    phase: str
    index: int
    magnitude: float


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 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 labels_to_target(labels: Sequence[tuple[str, str, int]], length: int, sigma_samples: float = 10.0) -> np.ndarray:
    target = np.zeros((5, length), dtype=DTYPE)
    x = np.arange(length, dtype=DTYPE)
    for phase, _group, idx in labels:
        pulse = np.exp(-0.5 * ((x - idx) / sigma_samples) ** 2)
        target[PHASE_TO_CHANNEL[phase]] = np.maximum(target[PHASE_TO_CHANNEL[phase]], pulse.astype(DTYPE))
    target[0] = np.clip(1.0 - target[1:].sum(axis=0), 0.0, 1.0)
    return target


def load_split_records(records_path: Path, h5_path: Path) -> tuple[dict[str, RecordInfo], list[PhaseExample]]:
    payload = json.loads(records_path.read_text())
    raw_records = payload["records"]
    event_ids = sorted({row["event"] for row in raw_records})
    with h5py.File(h5_path, "r") as h5:
        mag_by_event = {event: float(h5[event].attrs.get("event_magnitude", float("nan"))) for event in event_ids}

    records: dict[str, RecordInfo] = {}
    examples: list[PhaseExample] = []
    for row in raw_records:
        picks = tuple(
            ManualPick(phase=p["phase"], index=int(p["index"]))
            for p in row["phases"]
            if p.get("source") == "MANUAL"
        )
        if not picks:
            continue
        mag = mag_by_event.get(row["event"], float("nan"))
        if not math.isfinite(mag):
            continue
        record_id = f"{row['event']}/{row['station']}"
        records[record_id] = RecordInfo(
            event=row["event"],
            station=row["station"],
            length=int(row["length"]),
            magnitude=mag,
            picks=picks,
        )
        examples.extend(PhaseExample(record_id, pick.phase, pick.index, mag) for pick in picks)
    return records, examples


def window_start(record: RecordInfo, anchor_index: int, length: int, padlen: int, rng: np.random.Generator | None) -> int:
    if record.length <= length:
        return 0
    if rng is None:
        offset = length // 2
    else:
        low = min(padlen, length - 1)
        high = max(low + 1, length - padlen)
        offset = int(rng.integers(low, high))
    return int(np.clip(anchor_index - offset, 0, record.length - length))


def load_window(
    h5,
    records: dict[str, RecordInfo],
    example: PhaseExample,
    length: int,
    padlen: int,
    rng: np.random.Generator | None,
) -> tuple[np.ndarray, np.ndarray, int]:
    record = records[example.record_id]
    station = h5[record.event][record.station]
    comps = component_keys(station)
    if comps is None:
        raise RuntimeError(f"Missing components for {example.record_id}")
    start = window_start(record, example.index, length, padlen, rng)
    stop = min(start + length, record.length)
    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
    labels = []
    for pick in record.picks:
        rel = pick.index - start
        if 0 <= rel < length:
            labels.append((pick.phase, PHASE_TO_GROUP[pick.phase], rel))
    return normalize_wave(wave), labels_to_target(labels, length), example.index - start


def materialize_batch(
    h5,
    records: dict[str, RecordInfo],
    pool: Sequence[PhaseExample],
    selected: np.ndarray,
    length: int,
    padlen: int,
    rng: np.random.Generator | None,
) -> tuple[np.ndarray, np.ndarray]:
    waves = np.zeros((len(selected), length, 3), dtype=DTYPE)
    targets = np.zeros((len(selected), 5, length), dtype=DTYPE)
    for i, idx in enumerate(selected):
        wave, target, _rel = load_window(h5, records, pool[int(idx)], length, padlen, rng)
        waves[i] = wave
        targets[i] = target
    return waves, targets


def choose_matched_pools(
    train_examples: list[PhaseExample],
    magnitude_threshold: float,
    seed: int,
) -> tuple[list[PhaseExample], list[PhaseExample]]:
    high = [ex for ex in train_examples if ex.magnitude >= magnitude_threshold]
    if not high:
        raise RuntimeError("No high-magnitude manual phase examples found.")
    ordered_all = sorted(train_examples, key=lambda ex: (ex.record_id, ex.phase, ex.index))
    rng = np.random.default_rng(seed)
    all_idx = np.sort(rng.choice(len(ordered_all), size=len(high), replace=False))
    matched_all = [ordered_all[int(i)] for i in all_idx]
    return matched_all, high


def train_model(
    h5_path: Path,
    records: dict[str, RecordInfo],
    pool: list[PhaseExample],
    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)
    log_csv.parent.mkdir(parents=True, exist_ok=True)
    rng = np.random.default_rng(seed)
    with h5py.File(h5_path, "r") as h5, log_csv.open("w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["step", "loss"])
        for step in range(steps):
            selected = rng.integers(0, len(pool), size=batch_size)
            waves, targets = materialize_batch(h5, records, pool, selected, length, padlen, rng)
            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 % 25 == 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 empty_counts() -> dict:
    return {
        key: {"n": 0, "tp": 0}
        for key in ("P", "S", "combined", "Pg", "Sg", "Pn", "Sn")
    }


def evaluate_recall(
    h5_path: Path,
    records: dict[str, RecordInfo],
    examples: list[PhaseExample],
    model: BRNN,
    length: int,
    padlen: int,
    threshold: float,
    tolerance: int,
    min_sep: int,
    batch_size: int,
    device: torch.device,
) -> dict:
    counts = empty_counts()
    model.eval()
    with h5py.File(h5_path, "r") as h5, torch.no_grad():
        for start in range(0, len(examples), batch_size):
            batch_examples = examples[start : start + batch_size]
            waves = np.zeros((len(batch_examples), length, 3), dtype=DTYPE)
            rel_indices: list[int] = []
            for i, example in enumerate(batch_examples):
                wave, _target, rel = load_window(h5, records, example, length, padlen, rng=None)
                waves[i] = wave
                rel_indices.append(rel)
            xb = torch.from_numpy(waves).to(device).permute(0, 2, 1)
            out = model(xb).detach().cpu().numpy()
            for i, example in enumerate(batch_examples):
                phase = example.phase
                group = PHASE_TO_GROUP[phase]
                prob = out[i, GROUP_TO_CHANNELS[group], :].max(axis=0)
                pred = [idx for idx, _score in find_peaks(prob, threshold, min_sep)]
                hit = any(abs(idx - rel_indices[i]) <= tolerance for idx in pred)
                for key in (phase, group, "combined"):
                    counts[key]["n"] += 1
                    counts[key]["tp"] += int(hit)
            if start and start % (batch_size * 20) == 0:
                print(f"evaluated {start:,}/{len(examples):,} examples", flush=True)
    for key, row in counts.items():
        row["recall"] = row["tp"] / row["n"] if row["n"] else None
    return counts


def summarize_pool(examples: Sequence[PhaseExample]) -> dict:
    out = {"total_phase_examples": len(examples), "by_phase": {}, "by_group": {}, "event_count": 0, "record_count": 0}
    events = set()
    records = set()
    for ex in examples:
        event = ex.record_id.split("/", 1)[0]
        events.add(event)
        records.add(ex.record_id)
        out["by_phase"][ex.phase] = out["by_phase"].get(ex.phase, 0) + 1
        group = PHASE_TO_GROUP[ex.phase]
        out["by_group"][group] = out["by_group"].get(group, 0) + 1
    out["event_count"] = len(events)
    out["record_count"] = len(records)
    return out


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--h5", type=Path, default=ROOT / "data" / "credit-x1.h5")
    parser.add_argument("--train-records", type=Path, default=ROOT / "outputs" / "snr_transfer_seed20260609" / "records_train_all.json")
    parser.add_argument("--test-records", type=Path, default=ROOT / "outputs" / "snr_transfer_seed20260609" / "records_test_all.json")
    parser.add_argument("--base-ckpt", type=Path, default=ROOT / "ckpt" / "pnsn.v3.pt")
    parser.add_argument("--out-dir", type=Path, default=ROOT / "outputs" / "mag25_manual_phase_transfer")
    parser.add_argument("--magnitude-threshold", type=float, default=2.5)
    parser.add_argument("--seed", type=int, default=20260628)
    parser.add_argument("--train-steps", type=int, default=500)
    parser.add_argument("--train-batch", type=int, default=16)
    parser.add_argument("--eval-batch", type=int, default=64)
    parser.add_argument("--length", type=int, default=5120)
    parser.add_argument("--padlen", type=int, default=512)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--threshold", type=float, default=0.1)
    parser.add_argument("--tolerance-samples", type=int, default=100)
    parser.add_argument("--min-sep", type=int, default=50)
    parser.add_argument("--eval-max-examples", type=int, default=0)
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--resume", action="store_true")
    args = parser.parse_args()

    device = torch.device(args.device)
    args.out_dir.mkdir(parents=True, exist_ok=True)
    train_records, train_examples = load_split_records(args.train_records, args.h5)
    test_records, test_examples = load_split_records(args.test_records, args.h5)
    matched_all, high_train = choose_matched_pools(train_examples, args.magnitude_threshold, args.seed)
    high_eval = [ex for ex in test_examples if ex.magnitude >= args.magnitude_threshold]
    if args.eval_max_examples and args.eval_max_examples < len(high_eval):
        ordered = sorted(high_eval, key=lambda ex: (ex.record_id, ex.phase, ex.index))
        rng = np.random.default_rng(args.seed + 17)
        idx = np.sort(rng.choice(len(ordered), size=args.eval_max_examples, replace=False))
        high_eval = [ordered[int(i)] for i in idx]

    pools = {
        "allmag_matched": matched_all,
        "mag25plus": high_train,
    }
    models = {}
    for i, (name, pool) in enumerate(pools.items()):
        models[name] = train_model(
            h5_path=args.h5,
            records=train_records,
            pool=pool,
            base_ckpt=args.base_ckpt,
            out_ckpt=args.out_dir / f"pnsn.v3.transfer.{name}.pt",
            log_csv=args.out_dir / f"train_log_{name}.csv",
            seed=args.seed + i * 1009,
            steps=args.train_steps,
            batch_size=args.train_batch,
            length=args.length,
            padlen=args.padlen,
            lr=args.lr,
            device=device,
            resume=args.resume,
        )

    eval_results = {}
    for name, model in models.items():
        eval_results[name] = evaluate_recall(
            h5_path=args.h5,
            records=test_records,
            examples=high_eval,
            model=model,
            length=args.length,
            padlen=args.padlen,
            threshold=args.threshold,
            tolerance=args.tolerance_samples,
            min_sep=args.min_sep,
            batch_size=args.eval_batch,
            device=device,
        )

    delta = {}
    for key in eval_results["allmag_matched"]:
        a = eval_results["allmag_matched"][key]["recall"]
        b = eval_results["mag25plus"][key]["recall"]
        delta[key] = None if a is None or b is None else b - a
    summary = {
        "inputs": {
            "h5": str(args.h5),
            "train_records": str(args.train_records),
            "test_records": str(args.test_records),
            "base_ckpt": str(args.base_ckpt),
            "magnitude_threshold": args.magnitude_threshold,
            "seed": args.seed,
            "train_steps": args.train_steps,
            "train_batch": args.train_batch,
            "eval_batch": args.eval_batch,
            "length": args.length,
            "threshold": args.threshold,
            "tolerance_samples": args.tolerance_samples,
            "eval_max_examples": args.eval_max_examples,
        },
        "pool_summary": {
            "all_train_manual_phase_examples": summarize_pool(train_examples),
            "allmag_matched_train_pool": summarize_pool(matched_all),
            "mag25plus_train_pool": summarize_pool(high_train),
            "mag25plus_eval_pool": summarize_pool(high_eval),
        },
        "eval_manual_pick_recall_on_mag25plus": eval_results,
        "delta_mag25plus_minus_allmag": delta,
    }
    summary_path = args.out_dir / "summary.json"
    summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")

    rows_path = args.out_dir / "recall_table.csv"
    with rows_path.open("w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["phase_or_group", "allmag_recall", "mag25plus_recall", "delta", "n_eval"])
        for key in ("combined", "P", "S", "Pg", "Sg", "Pn", "Sn"):
            writer.writerow(
                [
                    key,
                    eval_results["allmag_matched"][key]["recall"],
                    eval_results["mag25plus"][key]["recall"],
                    delta[key],
                    eval_results["allmag_matched"][key]["n"],
                ]
            )
    print(json.dumps(summary["pool_summary"], indent=2), flush=True)
    print(json.dumps(summary["eval_manual_pick_recall_on_mag25plus"], indent=2), flush=True)
    print(f"wrote {summary_path}")
    print(f"wrote {rows_path}")


if __name__ == "__main__":
    main()