#!/usr/bin/env python3
"""Reproduce paper statistics with a fixed seed.

The script uses CREDIT-X1Local records to build deterministic single-event and
two-event mixture windows, fine-tunes the pretrained Pn/Sn picker, and exports
metrics plus publication-ready composite figures.
"""

from __future__ import annotations

import argparse
import csv
import datetime as dt
import json
import math
import os
import random
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, List, Sequence, Tuple

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


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]}
COMPONENTS = ("BHE", "BHN", "BHZ")
DTYPE = np.float32


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


@dataclass(frozen=True)
class Record:
    event: str
    station: str
    length: int
    delta: float
    distance_km: float
    phases: Tuple[PhasePick, ...]


@dataclass(frozen=True)
class CropSpec:
    rec_idx: int
    start: int
    amp: float


@dataclass(frozen=True)
class SampleSpec:
    crops: Tuple[CropSpec, ...]
    kind: str


def parse_time(value: str) -> dt.datetime:
    return dt.datetime.strptime(value, "%Y/%m/%d %H:%M:%S.%f")


def choose_phase(station, phase: str, prefer_manual: bool = True) -> Tuple[str, str] | None:
    manual_key = f"MANUAL.TRAVTIME.{phase}"
    rnn_key = f"RNN.TRAVTIME.{phase}"
    if prefer_manual and manual_key in station.attrs:
        value = station.attrs[manual_key]
        if isinstance(value, str) and "/" in value:
            return value, "MANUAL"
    if rnn_key in station.attrs:
        value = station.attrs[rnn_key]
        tag = station.attrs.get(f"{rnn_key}.tag", "")
        if isinstance(value, str) and "/" in value and tag == "Y":
            return value, "RNN.tagY"
    if not prefer_manual and manual_key in station.attrs:
        value = station.attrs[manual_key]
        if isinstance(value, str) and "/" in value:
            return value, "MANUAL"
    return 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 = {}
    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 build_records(
    h5_path: Path,
    key_path: Path,
    split: str,
    max_events: int | None,
    prefer_manual: bool = True,
) -> List[Record]:
    keys = np.load(key_path)[split]
    if max_events is not None:
        keys = keys[:max_events]
    records: List[Record] = []
    with h5py.File(h5_path, "r") as h5:
        for event_key in keys:
            event_key = str(event_key)
            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)
                lengths = [int(station[c].shape[0]) for c in comps]
                length = min(lengths)
                picks: List[PhasePick] = []
                for phase in ("Pg", "Sg", "Pn", "Sn"):
                    chosen = choose_phase(station, phase, prefer_manual=prefer_manual)
                    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),
                    )
                )
    return records


def crop_start_for_record(record: Record, rng: np.random.Generator, length: int, padlen: int) -> int | None:
    if record.length < length:
        return None
    phase_indices = np.array([p.index for p in record.phases], dtype=np.int64)
    for _ in range(24):
        anchor = int(rng.choice(phase_indices))
        offset = int(rng.integers(padlen, max(padlen + 1, length - padlen)))
        start = anchor - offset
        if 0 <= start <= record.length - length:
            return start
    anchor = int(rng.choice(phase_indices))
    return int(np.clip(anchor - length // 2, 0, record.length - length))


def make_specs(
    records: Sequence[Record],
    n_samples: int,
    seed: int,
    length: int,
    padlen: int,
    double_prob: float = 0.5,
    valid_indices: Sequence[int] | None = None,
) -> List[SampleSpec]:
    valid = list(valid_indices) if valid_indices is not None else [i for i, r in enumerate(records) if r.length >= length]
    if not valid:
        raise RuntimeError("No records are long enough for the requested window length.")
    rng = np.random.default_rng(seed)
    specs: List[SampleSpec] = []
    while len(specs) < n_samples:
        is_double = rng.random() < double_prob
        crop_count = 2 if is_double else 1
        crops: List[CropSpec] = []
        for j in range(crop_count):
            rec_idx = int(rng.choice(valid))
            start = crop_start_for_record(records[rec_idx], rng, length, padlen)
            if start is None:
                crops = []
                break
            amp = 1.0 if j == 0 else float(rng.uniform(0.2, 5.0))
            crops.append(CropSpec(rec_idx, start, amp))
        if crops:
            specs.append(SampleSpec(tuple(crops), "double" if is_double else "single"))
    return specs


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 load_crop(h5, record: Record, crop: CropSpec, length: int) -> Tuple[np.ndarray, List[Tuple[str, str, int]]]:
    station = h5[record.event][record.station]
    comps = component_keys(station)
    if comps is None:
        raise RuntimeError(f"Missing components for {record.event}/{record.station}")
    data = [station[c][crop.start : crop.start + length] for c in comps]
    wave = np.stack(data, axis=1)
    wave = normalize_wave(wave) * crop.amp
    labels: List[Tuple[str, str, int]] = []
    for pick in record.phases:
        rel = pick.index - crop.start
        if 0 <= rel < length:
            labels.append((pick.phase, PHASE_TO_GROUP[pick.phase], rel))
    return wave, labels


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, _, idx in labels:
        ch = PHASE_TO_CHANNEL[phase]
        pulse = np.exp(-0.5 * ((x - idx) / sigma_samples) ** 2)
        target[ch] = np.maximum(target[ch], pulse.astype(DTYPE))
    target[0] = np.clip(1.0 - target[1:].sum(axis=0), 0.0, 1.0)
    return target


def materialize_samples(
    h5_path: Path,
    records: Sequence[Record],
    specs: Sequence[SampleSpec],
    length: int,
    progress_interval: int = 0,
) -> Tuple[np.ndarray, List[List[Tuple[str, str, int]]], np.ndarray, List[str]]:
    waves = np.zeros((len(specs), length, 3), dtype=DTYPE)
    labels_all: List[List[Tuple[str, str, int]]] = []
    targets = np.zeros((len(specs), 5, length), dtype=DTYPE)
    kinds: List[str] = []
    with h5py.File(h5_path, "r") as h5:
        for i, spec in enumerate(specs):
            if progress_interval > 0 and (i % progress_interval == 0 or i == len(specs) - 1):
                print(f"materialized eval sample {i:,}/{len(specs):,}", flush=True)
            mixed = np.zeros((length, 3), dtype=DTYPE)
            labels: List[Tuple[str, str, int]] = []
            for crop in spec.crops:
                wave, crop_labels = load_crop(h5, records[crop.rec_idx], crop, length)
                mixed += wave
                labels.extend(crop_labels)
            mixed = normalize_wave(mixed)
            waves[i] = mixed
            labels_all.append(labels)
            targets[i] = labels_to_target(labels, length)
            kinds.append(spec.kind)
    return waves, labels_all, targets, kinds


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 find_peaks(prob: np.ndarray, threshold: float, min_sep: int) -> List[Tuple[int, float]]:
    if len(prob) < 3:
        return []
    mask = (prob[1:-1] >= threshold) & (prob[1:-1] >= prob[:-2]) & (prob[1:-1] > prob[2:])
    candidates = np.where(mask)[0] + 1
    if candidates.size == 0:
        return []
    order = candidates[np.argsort(prob[candidates])[::-1]]
    selected: List[int] = []
    for idx in order:
        if all(abs(int(idx) - old) >= min_sep for old in selected):
            selected.append(int(idx))
    selected.sort()
    return [(idx, float(prob[idx])) for idx in selected]


def match_predictions(
    true_indices: Sequence[int],
    pred_indices: Sequence[int],
    tolerance: int,
) -> Tuple[int, int, int, List[float]]:
    unmatched = set(range(len(true_indices)))
    tp = 0
    fp = 0
    errors: List[float] = []
    for pred in pred_indices:
        best = None
        best_dist = tolerance + 1
        for ti in unmatched:
            dist = abs(pred - true_indices[ti])
            if dist <= tolerance and dist < best_dist:
                best = ti
                best_dist = dist
        if best is None:
            fp += 1
        else:
            unmatched.remove(best)
            tp += 1
            errors.append((pred - true_indices[best]) * 0.01)
    fn = len(unmatched)
    return tp, fp, fn, errors


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):
            batch = torch.from_numpy(waves[start : start + batch_size]).to(device)
            batch = batch.permute(0, 2, 1)
            out = model(batch).detach().cpu().numpy()
            outputs.append(out)
    return np.concatenate(outputs, axis=0)


def evaluate_outputs(
    outputs: np.ndarray,
    labels_all: Sequence[Sequence[Tuple[str, str, int]]],
    kinds: Sequence[str],
    thresholds: Sequence[float],
    min_sep: int,
    tolerance: int,
) -> Dict:
    result = {"thresholds": list(map(float, thresholds)), "all": {}, "double": {}}
    for subset_name, subset_mask in {
        "all": np.ones(len(labels_all), dtype=bool),
        "double": np.array([k == "double" for k in kinds], dtype=bool),
    }.items():
        subset_result = {}
        for group in ("P", "S"):
            rows = []
            for thr in thresholds:
                tp = fp = fn = 0
                errors: List[float] = []
                for i, labels in enumerate(labels_all):
                    if not subset_mask[i]:
                        continue
                    true_idx = [idx for _, g, idx in labels if g == group]
                    prob = outputs[i, GROUP_TO_CHANNELS[group], :].max(axis=0)
                    pred_idx = [idx for idx, _ in find_peaks(prob, thr, min_sep)]
                    mtp, mfp, mfn, merr = match_predictions(true_idx, pred_idx, tolerance)
                    tp += mtp
                    fp += mfp
                    fn += mfn
                    errors.extend(merr)
                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
                row = {
                    "threshold": float(thr),
                    "tp": int(tp),
                    "fp": int(fp),
                    "fn": int(fn),
                    "precision": float(precision),
                    "recall": float(recall),
                    "f1": float(f1),
                    "mean_error_s": float(np.mean(errors)) if errors else None,
                    "std_error_s": float(np.std(errors)) if errors else None,
                    "errors_s": errors if abs(thr - 0.1) < 1e-9 else [],
                }
                rows.append(row)
            subset_result[group] = rows
        result[subset_name] = subset_result
    return result


def metric_at(metrics: Dict, subset: str, group: str, threshold: float) -> Dict:
    rows = metrics[subset][group]
    return min(rows, key=lambda r: abs(r["threshold"] - threshold))


def train_transfer(
    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,
) -> BRNN:
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    model = BRNN().to(device)
    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)
    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 % 25 == 0 or step == steps - 1:
                print(f"train step {step:05d}/{steps} loss={loss_value:.3f}", flush=True)
    torch.save(model.state_dict(), out_ckpt)
    return model


def load_model(ckpt: Path, device: torch.device) -> BRNN:
    model = BRNN().to(device)
    model.load_state_dict(torch.load(ckpt, map_location="cpu"))
    model.eval()
    return model


def plot_training_samples(waves: np.ndarray, labels_all, out: Path) -> None:
    single_idx = 0
    double_idx = next((i for i, labels in enumerate(labels_all) if len(labels) >= 3), min(1, len(labels_all) - 1))
    fig, axes = plt.subplots(2, 2, figsize=(11, 6.8), dpi=220)
    time = np.arange(waves.shape[1]) * 0.01
    for ax, idx, title in [
        (axes[0, 0], single_idx, "(a) Single-event waveform"),
        (axes[0, 1], double_idx, "(b) Two-event mixed waveform"),
    ]:
        for ci, lab in enumerate(["E", "N", "Z"]):
            ax.plot(time, waves[idx, :, ci] + (2 - ci) * 2.2, lw=0.55, color="black")
            ax.text(time[0] - 1.0, (2 - ci) * 2.2, lab, va="center", ha="right", fontsize=8)
        ax.set_title(title, loc="left", fontsize=10)
        ax.set_xlabel("Time (s)")
        ax.set_yticks([])
        ax.set_xlim(time[0], time[-1])
    for ax, idx, title in [
        (axes[1, 0], single_idx, "(c) Single-event labels"),
        (axes[1, 1], double_idx, "(d) Mixed-event labels"),
    ]:
        target = labels_to_target(labels_all[idx], waves.shape[1])
        ax.plot(time, target[1] + target[3], color="#d62728", label="P", lw=0.9)
        ax.plot(time, target[2] + target[4], color="#1f77b4", label="S", lw=0.9)
        ax.set_title(title, loc="left", fontsize=10)
        ax.set_xlabel("Time (s)")
        ax.set_ylabel("Probability")
        ax.set_ylim(-0.05, 1.05)
        ax.set_xlim(time[0], time[-1])
        ax.legend(frameon=False, fontsize=8, loc="upper right")
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)


def plot_pr(base_metrics: Dict, multi_metrics: Dict, out: Path) -> None:
    fig, axes = plt.subplots(2, 2, figsize=(9.2, 6.8), dpi=220, sharex=True, sharey=True)
    panels = [
        (axes[0, 0], base_metrics, "P", "(a) RNN, P"),
        (axes[0, 1], base_metrics, "S", "(b) RNN, S"),
        (axes[1, 0], multi_metrics, "P", "(c) Multi-RNN, P"),
        (axes[1, 1], multi_metrics, "S", "(d) Multi-RNN, S"),
    ]
    for ax, metrics, group, title in panels:
        rows = metrics["all"][group]
        x = [r["threshold"] for r in rows]
        ax.plot(x, [r["precision"] for r in rows], marker="o", ms=3, label="Precision")
        ax.plot(x, [r["recall"] for r in rows], marker="s", ms=3, label="Recall")
        ax.plot(x, [r["f1"] for r in rows], marker="^", ms=3, label="F1")
        ax.set_title(title, loc="left", fontsize=10)
        ax.grid(True, alpha=0.25)
        ax.set_ylim(0, 1.02)
        ax.set_xlabel("Minimum confidence threshold")
        ax.set_ylabel("Value")
    axes[0, 0].legend(frameon=False, fontsize=8, loc="best")
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)


def plot_errors(base_metrics: Dict, multi_metrics: Dict, subset: str, out: Path) -> None:
    fig, axes = plt.subplots(2, 2, figsize=(9.2, 6.8), dpi=220, sharex=True)
    panels = [
        (axes[0, 0], base_metrics, "P", "(a) RNN, P", "#ff6b6b"),
        (axes[0, 1], base_metrics, "S", "(b) RNN, S", "#6b6bff"),
        (axes[1, 0], multi_metrics, "P", "(c) Multi-RNN, P", "#ff6b6b"),
        (axes[1, 1], multi_metrics, "S", "(d) Multi-RNN, S", "#6b6bff"),
    ]
    for ax, metrics, group, title, color in panels:
        row = metric_at(metrics, subset, group, 0.1)
        errors = np.array(row["errors_s"], dtype=float)
        ax.hist(errors, bins=np.linspace(-2, 2, 81), color=color, alpha=0.82)
        ax.axvline(0, color="#2ca6df", ls="--", lw=0.9)
        text = (
            f"P={row['precision']:.3f}\n"
            f"R={row['recall']:.3f}\n"
            f"F1={row['f1']:.3f}\n"
            f"mean={row['mean_error_s'] * 1000:.1f} ms\n"
            f"std={row['std_error_s'] * 1000:.1f} ms"
        )
        ax.text(0.03, 0.95, text, transform=ax.transAxes, va="top", ha="left", fontsize=8)
        ax.set_title(title, loc="left", fontsize=10)
        ax.set_xlabel("Error (s)")
        ax.set_ylabel("Count")
        ax.set_xlim(-2, 2)
        ax.grid(True, alpha=0.18)
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)


def select_continuous_event(records: Sequence[Record], min_records: int = 24) -> str:
    by_event: Dict[str, List[Record]] = {}
    for record in records:
        if math.isnan(record.distance_km):
            continue
        by_event.setdefault(record.event, []).append(record)
    candidates = sorted(
        ((event, recs) for event, recs in by_event.items() if len(recs) >= min_records),
        key=lambda item: len(item[1]),
        reverse=True,
    )
    if not candidates:
        return max(by_event.items(), key=lambda item: len(item[1]))[0]
    return candidates[0][0]


def continuous_specs_for_event(records: Sequence[Record], event: str, length: int) -> List[Tuple[int, int]]:
    out = []
    for idx, record in enumerate(records):
        if record.event != event or record.length < length:
            continue
        phase_indices = [p.index for p in record.phases]
        if not phase_indices:
            continue
        start = int(np.clip(min(phase_indices) - 800, 0, record.length - length))
        out.append((idx, start))
    out.sort(key=lambda item: records[item[0]].distance_km)
    return out[:60]


def evaluate_continuous(
    model: BRNN,
    h5_path: Path,
    records: Sequence[Record],
    specs: Sequence[Tuple[int, int]],
    length: int,
    device: torch.device,
    threshold: float,
) -> Tuple[np.ndarray, List[List[Tuple[str, str, int]]], np.ndarray, Dict]:
    sample_specs = [SampleSpec((CropSpec(rec_idx, start, 1.0),), "continuous") for rec_idx, start in specs]
    waves, labels_all, _, _ = materialize_samples(h5_path, records, sample_specs, length)
    outputs = run_model(model, waves, device, batch_size=32)
    metrics = evaluate_outputs(
        outputs,
        labels_all,
        ["continuous"] * len(labels_all),
        [threshold],
        min_sep=50,
        tolerance=100,
    )
    return waves, labels_all, outputs, metrics


def plot_continuous(
    records: Sequence[Record],
    specs: Sequence[Tuple[int, int]],
    waves: np.ndarray,
    base_outputs: np.ndarray,
    multi_outputs: np.ndarray,
    threshold: float,
    out: Path,
) -> None:
    fig, axes = plt.subplots(1, 2, figsize=(9.2, 7.2), dpi=220, sharex=True, sharey=True)
    time = np.arange(waves.shape[1]) * 0.01
    for ax, outputs, title in [
        (axes[0], base_outputs, "(a) RNN"),
        (axes[1], multi_outputs, "(b) Multi-RNN"),
    ]:
        for i, (rec_idx, _) in enumerate(specs):
            rec = records[rec_idx]
            dist = rec.distance_km if not math.isnan(rec.distance_km) else i
            trace = waves[i, :, 2]
            trace = trace / (np.max(np.abs(trace)) + 1e-6) * 2.0 + dist
            ax.plot(time, trace, color="black", lw=0.35, alpha=0.75)
            for group, color in [("P", "#d62728"), ("S", "#1f77b4")]:
                prob = outputs[i, GROUP_TO_CHANNELS[group], :].max(axis=0)
                peaks = find_peaks(prob, threshold, min_sep=50)
                if peaks:
                    ax.scatter([p[0] * 0.01 for p in peaks], [dist] * len(peaks), s=6, color=color, alpha=0.9)
        ax.set_title(title, loc="left", fontsize=10)
        ax.set_xlabel("Time (s)")
        ax.grid(True, alpha=0.15)
    axes[0].set_ylabel("Epicentral distance (km)")
    axes[1].scatter([], [], s=12, color="#d62728", label="P")
    axes[1].scatter([], [], s=12, color="#1f77b4", label="S")
    axes[1].legend(frameon=False, fontsize=8, loc="upper right")
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)


def strip_errors(metrics: Dict) -> Dict:
    clean = json.loads(json.dumps(metrics))
    for subset in ("all", "double"):
        for group in ("P", "S"):
            for row in clean[subset][group]:
                row["error_count"] = len(row.get("errors_s", []))
                row.pop("errors_s", None)
    return clean


def summarize_records(records: Sequence[Record]) -> Dict:
    source_counts: Dict[str, int] = {}
    phase_counts: Dict[str, int] = {}
    for record in records:
        for pick in record.phases:
            source_counts[pick.source] = source_counts.get(pick.source, 0) + 1
            phase_counts[pick.phase] = phase_counts.get(pick.phase, 0) + 1
    return {
        "records": len(records),
        "source_counts": source_counts,
        "phase_counts": phase_counts,
    }


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/repro_seed20260607")
    parser.add_argument("--seed", type=int, default=20260607)
    parser.add_argument("--length", type=int, default=5120)
    parser.add_argument("--padlen", type=int, default=512)
    parser.add_argument("--train-steps", type=int, default=400)
    parser.add_argument("--train-batch", type=int, default=16)
    parser.add_argument("--eval-samples", type=int, default=20000)
    parser.add_argument("--eval-batch", type=int, default=64)
    parser.add_argument("--max-train-events", type=int, default=8000)
    parser.add_argument("--max-test-events", type=int, default=0)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--skip-train", action="store_true")
    args = parser.parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    h5_path = Path(args.h5)
    key_path = Path(args.keys)
    base_ckpt = Path(args.base_ckpt)
    multi_ckpt = out_dir / f"pnsn.v3.multirnn.seed{args.seed}.pt"

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    random.seed(args.seed)
    device = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
    print(f"device={device}", flush=True)

    max_test = None if args.max_test_events == 0 else args.max_test_events
    train_records = build_records(h5_path, key_path, "train", args.max_train_events)
    test_records = build_records(h5_path, key_path, "test", max_test)
    print(f"train records={len(train_records)} test records={len(test_records)}", flush=True)

    with (out_dir / "record_summary.json").open("w") as f:
        json.dump(
            {
                "train": summarize_records(train_records),
                "test": summarize_records(test_records),
                "seed": args.seed,
                "length_samples": args.length,
                "length_seconds": args.length * 0.01,
                "label_rule": "MANUAL.TRAVTIME preferred; RNN.TRAVTIME used only when tag=Y.",
            },
            f,
            ensure_ascii=False,
            indent=2,
        )

    if args.skip_train and multi_ckpt.exists():
        multi_model = load_model(multi_ckpt, device)
    else:
        multi_model = train_transfer(
            h5_path=h5_path,
            records=train_records,
            base_ckpt=base_ckpt,
            out_ckpt=multi_ckpt,
            log_csv=out_dir / "transfer_loss.csv",
            seed=args.seed,
            steps=args.train_steps,
            batch_size=args.train_batch,
            length=args.length,
            padlen=args.padlen,
            lr=args.lr,
            device=device,
        )

    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,
    )
    waves, labels_all, _, kinds = materialize_samples(h5_path, test_records, eval_specs, args.length)
    np.savez_compressed(
        out_dir / "eval_sample_summary.npz",
        kinds=np.array(kinds),
        label_counts=np.array([len(x) for x in labels_all], dtype=np.int16),
    )
    plot_training_samples(waves[:100], labels_all[:100], out_dir / "fig5_training_samples_composite.png")

    thresholds = [round(x, 1) for x in np.arange(0.1, 1.0, 0.1)]
    base_model = load_model(base_ckpt, device)
    base_outputs = run_model(base_model, waves, device, args.eval_batch)
    multi_outputs = run_model(multi_model, waves, device, args.eval_batch)
    base_metrics = evaluate_outputs(base_outputs, labels_all, kinds, thresholds, min_sep=50, tolerance=100)
    multi_metrics = evaluate_outputs(multi_outputs, labels_all, kinds, thresholds, min_sep=50, tolerance=100)

    plot_pr(base_metrics, multi_metrics, out_dir / "fig6_pr_composite.png")
    plot_errors(base_metrics, multi_metrics, "all", out_dir / "fig7_error_all_composite.png")
    plot_errors(base_metrics, multi_metrics, "double", out_dir / "fig8_error_double_composite.png")

    event = select_continuous_event(test_records)
    cont_specs = continuous_specs_for_event(test_records, event, min(args.length * 2, 10240))
    cont_len = min(args.length * 2, 10240)
    if len(cont_specs) >= 4:
        cont_base_waves, cont_labels, cont_base_outputs, cont_base_metrics = evaluate_continuous(
            base_model, h5_path, test_records, cont_specs, cont_len, device, threshold=0.1
        )
        _, _, cont_multi_outputs, cont_multi_metrics = evaluate_continuous(
            multi_model, h5_path, test_records, cont_specs, cont_len, device, threshold=0.1
        )
        plot_continuous(
            test_records,
            cont_specs,
            cont_base_waves,
            cont_base_outputs,
            cont_multi_outputs,
            0.1,
            out_dir / "fig9_continuous_composite.png",
        )
    else:
        cont_base_metrics = {}
        cont_multi_metrics = {}

    report = {
        "seed": args.seed,
        "device": str(device),
        "length_samples": args.length,
        "length_seconds": args.length * 0.01,
        "eval_samples": args.eval_samples,
        "eval_single_samples": int(sum(k == "single" for k in kinds)),
        "eval_double_samples": int(sum(k == "double" for k in kinds)),
        "threshold_for_error_stats": 0.1,
        "match_tolerance_seconds": 1.0,
        "min_peak_separation_seconds": 0.5,
        "base_metrics": strip_errors(base_metrics),
        "multi_metrics": strip_errors(multi_metrics),
        "continuous_event": event,
        "continuous_station_count": len(cont_specs),
        "continuous_base_metrics": strip_errors(cont_base_metrics) if cont_base_metrics else {},
        "continuous_multi_metrics": strip_errors(cont_multi_metrics) if cont_multi_metrics else {},
        "figures": {
            "fig5": str((out_dir / "fig5_training_samples_composite.png").resolve()),
            "fig6": str((out_dir / "fig6_pr_composite.png").resolve()),
            "fig7": str((out_dir / "fig7_error_all_composite.png").resolve()),
            "fig8": str((out_dir / "fig8_error_double_composite.png").resolve()),
            "fig9": str((out_dir / "fig9_continuous_composite.png").resolve()),
        },
    }
    with (out_dir / "metrics_report.json").open("w") as f:
        json.dump(report, f, ensure_ascii=False, indent=2)

    print(json.dumps(report, ensure_ascii=False, indent=2)[:4000], flush=True)


if __name__ == "__main__":
    main()