code/scripts/snr_transfer_phase_balanced_experiment.py

#!/usr/bin/env python3
"""Fine-tune the Pn/Sn picker with phase-aware SNR-filtered labels.

This variant keeps the original P thresholds (5 and 10 dB) but chooses lower
S thresholds so that the retained P- and S-label counts are approximately
balanced in the medium- and high-SNR training pools. Unlike the original
record-level SNR experiment, filtering is applied at the phase-label level:
labels below their phase-specific threshold are removed from the training
target, and records with no retained labels are excluded.

For stricter P/S joint training, use ``--filter-mode record-complete`` with
``--s-threshold-mode same-as-p``. That mode keeps a training record only when
every original P/S label in the record has finite SNR and passes its threshold,
and it keeps the full original label set in the target. Add
``--match-mode phase-composition`` when P-only, P+S, and S-only waveform counts
must be identical across the full, medium-SNR, and high-SNR training pools.

Use ``--filter-mode record-any`` when a less restrictive record-level filter is
desired: if any original P/S label in the record passes its phase-specific SNR
threshold, the whole waveform record is kept with all original labels.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import random
import sys
from collections import Counter
from pathlib import Path
from typing import Any

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 (
    COMPONENTS,
    DTYPE,
    PHASE_TO_GROUP,
    PhasePick,
    Record,
    evaluate_outputs,
    labels_to_target,
    load_crop,
    make_specs,
    materialize_samples,
    normalize_wave,
    run_model,
    strip_errors,
)
from scripts.snr_transfer_experiment import (
    build_records_sequential,
    component_keys,
    matched_records,
    metric_row,
    pick_snr_db,
    plot_summary,
    sample_batch,
    write_metrics_table,
)


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


H5AccessIndex = dict[str, tuple[str, str, str]]


def build_h5_access_index(h5_path: Path, records: list[Record], cache_path: Path) -> H5AccessIndex:
    if cache_path.exists():
        raw = json.loads(cache_path.read_text(encoding="utf-8"))
        paths = raw.get("paths", raw)
        print(f"loaded HDF5 access index {cache_path} entries={len(paths):,}", flush=True)
        return {str(k): tuple(v) for k, v in paths.items()}

    print(f"building HDF5 access index {cache_path} from local record keys", flush=True)
    paths: dict[str, list[str]] = {
        record_key(record): [f"/{record.event}/{record.station}/{component}" for component in COMPONENTS]
        for record in records
    }
    cache_path.parent.mkdir(parents=True, exist_ok=True)
    tmp = cache_path.with_suffix(".tmp")
    tmp.write_text(
        json.dumps({"h5": str(h5_path), "record_count": len(records), "paths": paths}, ensure_ascii=False),
        encoding="utf-8",
    )
    tmp.replace(cache_path)
    return {k: tuple(v) for k, v in paths.items()}


def phase_snr_key(record: Record, pick: PhasePick, pick_index: int) -> str:
    return f"{record.event}/{record.station}/{pick_index}:{pick.phase}:{pick.index}:{pick.source}"


def phase_counts(records: list[Record]) -> dict[str, int]:
    counts = Counter()
    for record in records:
        for pick in record.phases:
            counts[PHASE_TO_GROUP[pick.phase]] += 1
    return {"P": int(counts["P"]), "S": int(counts["S"]), "total": int(counts["P"] + counts["S"])}


PHASE_COMPOSITION_ORDER = ("P_only", "PS", "S_only", "none")


def record_phase_composition(record: Record) -> str:
    groups = {PHASE_TO_GROUP[pick.phase] for pick in record.phases}
    if groups == {"P", "S"}:
        return "PS"
    if groups == {"P"}:
        return "P_only"
    if groups == {"S"}:
        return "S_only"
    return "none"


def phase_composition_counts(records: list[Record]) -> dict[str, int]:
    counts = Counter(record_phase_composition(record) for record in records)
    return {key: int(counts[key]) for key in PHASE_COMPOSITION_ORDER}


def matched_records_by_phase_composition(
    records: list[Record],
    targets: dict[str, int],
    seed: int,
) -> list[Record]:
    by_composition: dict[str, list[Record]] = {key: [] for key in PHASE_COMPOSITION_ORDER}
    for record in records:
        by_composition[record_phase_composition(record)].append(record)

    rng = np.random.default_rng(seed)
    selected: list[Record] = []
    for key in PHASE_COMPOSITION_ORDER:
        target = int(targets.get(key, 0))
        if target == 0:
            continue
        pool = sorted(by_composition[key], key=record_key)
        if len(pool) < target:
            raise RuntimeError(f"Cannot draw {target} {key} records from a pool of {len(pool)} records.")
        idx = np.sort(rng.choice(len(pool), size=target, replace=False))
        selected.extend(pool[int(i)] for i in idx)
    return sorted(selected, key=record_key)


def compute_phase_snr_cache(h5_path: Path, records: list[Record], cache_path: Path) -> dict[str, float]:
    if cache_path.exists():
        return {k: float(v) for k, v in json.loads(cache_path.read_text(encoding="utf-8")).items()}

    cache_path.parent.mkdir(parents=True, exist_ok=True)
    values: 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]
            for pick_index, pick in enumerate(record.phases):
                snr = pick_snr_db(waves, pick.phase, pick.index)
                if snr is not None and math.isfinite(snr):
                    values[phase_snr_key(record, pick, pick_index)] = float(snr)
            if i % 5000 == 0:
                print(f"computed phase SNR for {i:,}/{len(records):,} records", flush=True)

    tmp = cache_path.with_suffix(".tmp")
    tmp.write_text(json.dumps(values, ensure_ascii=False), encoding="utf-8")
    tmp.replace(cache_path)
    return values


def collect_phase_snr(records: list[Record], phase_snr: dict[str, float]) -> dict[str, list[float]]:
    by_group = {"P": [], "S": []}
    for record in records:
        for pick_index, pick in enumerate(record.phases):
            value = phase_snr.get(phase_snr_key(record, pick, pick_index))
            if value is None:
                continue
            by_group[PHASE_TO_GROUP[pick.phase]].append(float(value))
    return by_group


def choose_s_threshold(s_values: list[float], target_count: int) -> tuple[float, int]:
    if target_count <= 0 or not s_values:
        raise RuntimeError("Cannot choose S threshold without P target and S values.")
    ordered = sorted(float(v) for v in s_values if math.isfinite(v))
    if target_count >= len(ordered):
        return ordered[0] - 1e-6, len(ordered)
    # Count kept is values >= threshold. Use the boundary value giving the
    # nearest count; subtract a tiny epsilon so equality is included.
    threshold = ordered[-target_count] - 1e-6
    kept = sum(v >= threshold for v in ordered)
    return float(threshold), int(kept)


def filter_records_by_phase_thresholds(
    records: list[Record],
    phase_snr: dict[str, float],
    p_threshold: float | None,
    s_threshold: float | None,
) -> list[Record]:
    if p_threshold is None and s_threshold is None:
        return list(records)
    filtered: list[Record] = []
    for record in records:
        kept: list[PhasePick] = []
        for pick_index, pick in enumerate(record.phases):
            value = phase_snr.get(phase_snr_key(record, pick, pick_index))
            if value is None:
                continue
            group = PHASE_TO_GROUP[pick.phase]
            threshold = p_threshold if group == "P" else s_threshold
            if threshold is not None and value >= threshold:
                kept.append(pick)
        if kept:
            filtered.append(
                Record(
                    event=record.event,
                    station=record.station,
                    length=record.length,
                    delta=record.delta,
                    distance_km=record.distance_km,
                    phases=tuple(kept),
                )
            )
    return filtered


def filter_records_by_complete_thresholds(
    records: list[Record],
    phase_snr: dict[str, float],
    p_threshold: float | None,
    s_threshold: float | None,
) -> list[Record]:
    if p_threshold is None and s_threshold is None:
        return list(records)
    filtered: list[Record] = []
    for record in records:
        keep_record = True
        for pick_index, pick in enumerate(record.phases):
            value = phase_snr.get(phase_snr_key(record, pick, pick_index))
            if value is None or not math.isfinite(value):
                keep_record = False
                break
            group = PHASE_TO_GROUP[pick.phase]
            threshold = p_threshold if group == "P" else s_threshold
            if threshold is not None and value < threshold:
                keep_record = False
                break
        if keep_record:
            filtered.append(record)
    return filtered


def filter_records_by_any_thresholds(
    records: list[Record],
    phase_snr: dict[str, float],
    p_threshold: float | None,
    s_threshold: float | None,
) -> list[Record]:
    if p_threshold is None and s_threshold is None:
        return list(records)
    filtered: list[Record] = []
    for record in records:
        keep_record = False
        for pick_index, pick in enumerate(record.phases):
            value = phase_snr.get(phase_snr_key(record, pick, pick_index))
            if value is None or not math.isfinite(value):
                continue
            group = PHASE_TO_GROUP[pick.phase]
            threshold = p_threshold if group == "P" else s_threshold
            if threshold is not None and value >= threshold:
                keep_record = True
                break
        if keep_record:
            filtered.append(record)
    return filtered


def summarize_snr(values: list[float]) -> dict[str, float | int]:
    arr = np.asarray(values, dtype=float)
    arr = arr[np.isfinite(arr)]
    if arr.size == 0:
        return {"count": 0}
    return {
        "count": int(arr.size),
        "min": float(np.min(arr)),
        "p25": float(np.percentile(arr, 25)),
        "median": float(np.median(arr)),
        "p75": float(np.percentile(arr, 75)),
        "p95": float(np.percentile(arr, 95)),
        "max": float(np.max(arr)),
    }


def write_phase_balance_table(rows: list[dict[str, Any]], out: Path) -> None:
    with out.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(
            f,
            fieldnames=[
                "slug",
                "label",
                "p_threshold_db",
                "s_threshold_db",
                "candidate_records",
                "candidate_P_only_records",
                "candidate_PS_records",
                "candidate_S_only_records",
                "candidate_none_records",
                "candidate_P_labels",
                "candidate_S_labels",
                "matched_records",
                "matched_P_only_records",
                "matched_PS_records",
                "matched_S_only_records",
                "matched_none_records",
                "matched_P_labels",
                "matched_S_labels",
            ],
        )
        writer.writeheader()
        writer.writerows(rows)


def load_or_materialize_eval_samples(
    h5_path: Path,
    test_records: list[Record],
    access_index: H5AccessIndex,
    cache_path: Path,
    seed: int,
    n_samples: int,
    length: int,
    padlen: int,
) -> tuple[np.ndarray, list[list[tuple[str, str, int]]], list[str]]:
    if cache_path.exists():
        print(f"loading cached eval samples {cache_path}", flush=True)
        data = np.load(cache_path, allow_pickle=True)
        return data["waves"], data["labels"].tolist(), data["kinds"].tolist()

    print(f"materializing eval samples {cache_path}", flush=True)
    eval_specs = make_specs(
        test_records,
        n_samples=n_samples,
        seed=seed,
        length=length,
        padlen=padlen,
        double_prob=0.5,
    )
    waves = np.zeros((len(eval_specs), length, 3), dtype=DTYPE)
    labels: list[list[tuple[str, str, int]]] = []
    kinds: list[str] = []
    with h5py.File(h5_path, "r") as h5:
        for i, spec in enumerate(eval_specs):
            if i % 1000 == 0 or i == len(eval_specs) - 1:
                print(f"materialized eval sample {i:,}/{len(eval_specs):,}", flush=True)
            mixed = np.zeros((length, 3), dtype=DTYPE)
            sample_labels: list[tuple[str, str, int]] = []
            for crop in spec.crops:
                wave, crop_labels = load_crop_indexed(h5, test_records[crop.rec_idx], crop, length, access_index)
                mixed += wave
                sample_labels.extend(crop_labels)
            waves[i] = normalize_wave(mixed)
            labels.append(sample_labels)
            kinds.append(spec.kind)
    cache_path.parent.mkdir(parents=True, exist_ok=True)
    np.savez(
        cache_path,
        waves=waves,
        labels=np.asarray(labels, dtype=object),
        kinds=np.asarray(kinds, dtype=object),
    )
    print(f"wrote cached eval samples {cache_path}", flush=True)
    return waves, labels, kinds


def load_crop_indexed(
    h5: h5py.File,
    record: Record,
    crop: Any,
    length: int,
    access_index: H5AccessIndex,
) -> tuple[np.ndarray, list[tuple[str, str, int]]]:
    paths = access_index.get(record_key(record))
    if paths is None:
        return load_crop(h5, record, crop, length)
    try:
        data = [h5[path][crop.start : crop.start + length] for path in paths]
    except KeyError:
        return load_crop(h5, record, crop, length)
    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 materialize_specs_from_open_h5(
    h5: h5py.File,
    records: list[Record],
    specs: list,
    length: int,
    access_index: H5AccessIndex,
) -> tuple[np.ndarray, np.ndarray]:
    waves = np.zeros((len(specs), length, 3), dtype=DTYPE)
    targets = np.zeros((len(specs), 5, length), dtype=DTYPE)
    for i, spec in enumerate(specs):
        mixed = np.zeros((length, 3), dtype=DTYPE)
        labels: list[tuple[str, str, int]] = []
        for crop in spec.crops:
            wave, crop_labels = load_crop_indexed(h5, records[crop.rec_idx], crop, length, access_index)
            mixed += wave
            labels.extend(crop_labels)
        mixed = normalize_wave(mixed)
        waves[i] = mixed
        targets[i] = labels_to_target(labels, length)
    return waves, targets


def sample_batch_open_h5(
    h5: h5py.File,
    records: list[Record],
    access_index: H5AccessIndex,
    valid_indices: list[int],
    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,
        valid_indices=valid_indices,
    )
    return materialize_specs_from_open_h5(h5, records, specs, length, access_index)


def train_model(
    h5_path: Path,
    records: list[Record],
    access_index: H5AccessIndex,
    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,
    checkpoint_every_steps: int,
) -> BRNN:
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    model = BRNN().to(device)
    progress_ckpt = out_ckpt.with_name(f"{out_ckpt.stem}.progress.pt")
    train_state_path = out_ckpt.with_name(f"{out_ckpt.stem}.train_state.json")
    if resume and out_ckpt.exists():
        state = {}
        if train_state_path.exists():
            state = json.loads(train_state_path.read_text(encoding="utf-8"))
        if state.get("complete", True):
            model.load_state_dict(torch.load(out_ckpt, map_location="cpu"))
            model.eval()
            print(f"loaded existing checkpoint {out_ckpt}", flush=True)
            return model

    if init_mode == "finetune":
        model.load_state_dict(torch.load(base_ckpt, map_location="cpu"))
    elif init_mode == "scratch":
        pass
    else:
        raise ValueError(f"Unsupported init mode: {init_mode}")

    model.train()
    loss_fn = Loss().to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
    start_step = 0
    if resume and progress_ckpt.exists():
        progress = torch.load(progress_ckpt, map_location="cpu")
        if (
            progress.get("init_mode") == init_mode
            and progress.get("target_steps") == steps
            and progress.get("seed") == seed
        ):
            model.load_state_dict(progress["model_state"])
            opt.load_state_dict(progress["optimizer_state"])
            start_step = int(progress.get("completed_steps", 0))
            print(f"resuming partial checkpoint {progress_ckpt} at step {start_step}/{steps}", flush=True)
        else:
            print(f"ignoring incompatible partial checkpoint {progress_ckpt}", flush=True)

    out_ckpt.parent.mkdir(parents=True, exist_ok=True)
    log_csv.parent.mkdir(parents=True, exist_ok=True)
    log_mode = "a" if start_step > 0 and log_csv.exists() else "w"
    valid_indices = [i for i, record in enumerate(records) if record.length >= length]
    if not valid_indices:
        raise RuntimeError("No training records are long enough for the requested window length.")
    with log_csv.open(log_mode, newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        if log_mode == "w":
            writer.writerow(["step", "loss"])
        with h5py.File(h5_path, "r") as train_h5:
            for step in range(start_step, steps):
                waves, targets = sample_batch_open_h5(
                    train_h5,
                    records,
                    access_index,
                    valid_indices,
                    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,
                    )
                    f.flush()
                completed_steps = step + 1
                if checkpoint_every_steps > 0 and (
                    completed_steps % checkpoint_every_steps == 0 or completed_steps == steps
                ):
                    torch.save(
                        {
                            "model_state": model.state_dict(),
                            "optimizer_state": opt.state_dict(),
                            "completed_steps": completed_steps,
                            "target_steps": steps,
                            "init_mode": init_mode,
                            "seed": seed,
                        },
                        progress_ckpt,
                    )
                    train_state_path.write_text(
                        json.dumps(
                            {
                                "complete": False,
                                "completed_steps": completed_steps,
                                "target_steps": steps,
                                "init_mode": init_mode,
                                "seed": seed,
                                "progress_checkpoint": str(progress_ckpt),
                            },
                            ensure_ascii=False,
                            indent=2,
                        ),
                        encoding="utf-8",
                    )
    torch.save(model.state_dict(), out_ckpt)
    train_state_path.write_text(
        json.dumps(
            {
                "complete": True,
                "completed_steps": steps,
                "target_steps": steps,
                "init_mode": init_mode,
                "seed": seed,
                "checkpoint": str(out_ckpt),
            },
            ensure_ascii=False,
            indent=2,
        ),
        encoding="utf-8",
    )
    if progress_ckpt.exists():
        progress_ckpt.unlink()
    model.eval()
    return model


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    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_phase_balanced_seed20260609")
    parser.add_argument(
        "--cache-dir",
        default=None,
        help="Directory for reusable record and phase-SNR caches. Defaults to --out-dir.",
    )
    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(
        "--scratch-train-steps",
        type=int,
        default=0,
        help="Override --train-steps for scratch initialization. Defaults to --train-steps.",
    )
    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(
        "--checkpoint-every-steps",
        type=int,
        default=0,
        help="Save resumable model and optimizer state during training every N steps.",
    )
    parser.add_argument(
        "--init-modes",
        nargs="+",
        choices=["finetune", "scratch"],
        default=["finetune"],
        help="Initialization modes to run on the same matched training pools.",
    )
    parser.add_argument(
        "--force-retrain-init-modes",
        nargs="+",
        choices=["finetune", "scratch"],
        default=[],
        help="Ignore existing checkpoints for these initialization modes even when --resume is set.",
    )
    parser.add_argument(
        "--subset-slugs",
        nargs="+",
        choices=["full", "p5_s_bal", "p10_s_bal"],
        default=["full", "p5_s_bal", "p10_s_bal"],
        help="Training subsets to run. Candidate matching is still computed across all subsets.",
    )
    parser.add_argument(
        "--merge-existing-summary",
        action="store_true",
        help="Merge newly evaluated rows into an existing summary.json by slug.",
    )
    parser.add_argument(
        "--train-only",
        action="store_true",
        help="Train or resume selected models and skip evaluation/summary plotting.",
    )
    parser.add_argument(
        "--no-match-train-size",
        action="store_true",
        help="Use all phase-filtered candidate records; training steps remain equal.",
    )
    parser.add_argument(
        "--match-mode",
        choices=["record-count", "phase-composition"],
        default="record-count",
        help=(
            "record-count matches only the total record count; phase-composition also "
            "matches the number of P-only, P+S, and S-only waveform records."
        ),
    )
    parser.add_argument(
        "--filter-mode",
        choices=["phase-label", "record-complete", "record-any"],
        default="phase-label",
        help=(
            "phase-label removes only labels below threshold; record-complete keeps "
            "only records whose original labels all pass and then trains on the full label set; "
            "record-any keeps the full label set when any original label passes."
        ),
    )
    parser.add_argument(
        "--s-threshold-mode",
        choices=["balanced", "same-as-p"],
        default="balanced",
        help="Choose S thresholds by P/S label-count balancing or set S thresholds equal to the P thresholds.",
    )
    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)
    cache_dir = Path(args.cache_dir) if args.cache_dir else out_dir
    cache_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
    train_tag = "all" if max_train is None else str(max_train)
    test_tag = "all" if max_test is None else str(max_test)
    train_records = build_records_sequential(
        h5_path, key_path, "train", max_train, cache_dir / f"records_train_{train_tag}.json"
    )
    test_records = build_records_sequential(
        h5_path, key_path, "test", max_test, cache_dir / f"records_test_{test_tag}.json"
    )
    print(f"train records={len(train_records):,} test records={len(test_records):,}", flush=True)
    train_access_index = build_h5_access_index(
        h5_path,
        train_records,
        cache_dir / f"h5_access_index_train_{train_tag}.json",
    )
    test_access_index = build_h5_access_index(
        h5_path,
        test_records,
        cache_dir / f"h5_access_index_test_{test_tag}.json",
    )

    phase_snr = compute_phase_snr_cache(h5_path, train_records, cache_dir / "train_phase_snr_db.json")
    by_group = collect_phase_snr(train_records, phase_snr)
    p5_count = sum(v >= 5.0 for v in by_group["P"])
    p10_count = sum(v >= 10.0 for v in by_group["P"])
    if args.s_threshold_mode == "balanced":
        s5, s5_count = choose_s_threshold(by_group["S"], p5_count)
        s10, s10_count = choose_s_threshold(by_group["S"], p10_count)
    else:
        s5 = 5.0
        s10 = 10.0
        s5_count = sum(v >= s5 for v in by_group["S"])
        s10_count = sum(v >= s10 for v in by_group["S"])
    print(
        f"phase-aware thresholds ({args.s_threshold_mode}, {args.filter_mode}): "
        f"P>=5/S>={s5:.3f} gives P={p5_count:,} S={s5_count:,}; "
        f"P>=10/S>={s10:.3f} gives P={p10_count:,} S={s10_count:,}",
        flush=True,
    )

    p5_label = f"P>=5 dB, S>={s5:.2f} dB"
    p10_label = f"P>=10 dB, S>={s10:.2f} dB"
    if args.filter_mode == "record-complete":
        p5_label += " record-complete"
        p10_label += " record-complete"
    elif args.filter_mode == "record-any":
        p5_label += " record-any"
        p10_label += " record-any"
    subset_defs = [
        ("full", "Full", None, None),
        ("p5_s_bal", p5_label, 5.0, s5),
        ("p10_s_bal", p10_label, 10.0, s10),
    ]
    if args.filter_mode == "record-complete":
        filter_fn = filter_records_by_complete_thresholds
    elif args.filter_mode == "record-any":
        filter_fn = filter_records_by_any_thresholds
    else:
        filter_fn = filter_records_by_phase_thresholds
    candidate_records = {
        slug: filter_fn(train_records, phase_snr, p_thr, s_thr)
        for slug, _label, p_thr, s_thr in subset_defs
    }
    candidate_counts = {slug: len(records) for slug, records in candidate_records.items()}
    candidate_phase_counts = {slug: phase_counts(records) for slug, records in candidate_records.items()}
    candidate_phase_composition_counts = {
        slug: phase_composition_counts(records) for slug, records in candidate_records.items()
    }
    if args.no_match_train_size:
        train_pools = candidate_records
        matched_train_records = None
        matched_train_composition = None
    elif args.match_mode == "phase-composition":
        matched_train_composition = {
            key: min(counts[key] for counts in candidate_phase_composition_counts.values())
            for key in PHASE_COMPOSITION_ORDER
        }
        matched_train_records = sum(matched_train_composition.values())
        train_pools = {
            slug: matched_records_by_phase_composition(records, matched_train_composition, args.seed + i * 8191)
            for i, (slug, _label, _p_thr, _s_thr) in enumerate(subset_defs)
            for records in [candidate_records[slug]]
        }
    else:
        matched_train_records = min(candidate_counts.values())
        matched_train_composition = None
        train_pools = {
            slug: matched_records(records, matched_train_records, args.seed + i * 8191)
            for i, (slug, _label, _p_thr, _s_thr) in enumerate(subset_defs)
            for records in [candidate_records[slug]]
        }
    matched_phase_counts = {slug: phase_counts(records) for slug, records in train_pools.items()}
    matched_phase_composition_counts = {
        slug: phase_composition_counts(records) for slug, records in train_pools.items()
    }
    balance_rows = []
    for slug, label, p_thr, s_thr in subset_defs:
        candidate_composition = candidate_phase_composition_counts[slug]
        matched_composition = matched_phase_composition_counts[slug]
        balance_rows.append(
            {
                "slug": slug,
                "label": label,
                "p_threshold_db": "" if p_thr is None else p_thr,
                "s_threshold_db": "" if s_thr is None else s_thr,
                "candidate_records": candidate_counts[slug],
                "candidate_P_only_records": candidate_composition["P_only"],
                "candidate_PS_records": candidate_composition["PS"],
                "candidate_S_only_records": candidate_composition["S_only"],
                "candidate_none_records": candidate_composition["none"],
                "candidate_P_labels": candidate_phase_counts[slug]["P"],
                "candidate_S_labels": candidate_phase_counts[slug]["S"],
                "matched_records": len(train_pools[slug]),
                "matched_P_only_records": matched_composition["P_only"],
                "matched_PS_records": matched_composition["PS"],
                "matched_S_only_records": matched_composition["S_only"],
                "matched_none_records": matched_composition["none"],
                "matched_P_labels": matched_phase_counts[slug]["P"],
                "matched_S_labels": matched_phase_counts[slug]["S"],
            }
        )
    write_phase_balance_table(balance_rows, out_dir / "phase_balance_table.csv")
    print("candidate counts=" + json.dumps(candidate_counts, ensure_ascii=False), flush=True)
    print("phase balance=" + json.dumps(balance_rows, ensure_ascii=False), flush=True)

    eval_cache_path = (
        cache_dir
        / f"eval_seed{args.seed}_test{test_tag}_n{args.eval_samples}_len{args.length}_pad{args.padlen}.npz"
    )
    eval_data: tuple[np.ndarray, list[list[tuple[str, str, int]]], list[str]] | None = None

    def get_eval_data() -> tuple[np.ndarray, list[list[tuple[str, str, int]]], list[str]]:
        nonlocal eval_data
        if eval_data is None:
            eval_data = load_or_materialize_eval_samples(
                h5_path=h5_path,
                test_records=test_records,
                access_index=test_access_index,
                cache_path=eval_cache_path,
                seed=args.seed + 17,
                n_samples=args.eval_samples,
                length=args.length,
                padlen=args.padlen,
            )
        return eval_data

    thresholds = [round(x, 1) for x in np.arange(0.1, 1.0, 0.1)]
    train_steps_by_init_mode = {
        "finetune": args.train_steps,
        "scratch": args.scratch_train_steps if args.scratch_train_steps > 0 else args.train_steps,
    }
    force_retrain_init_modes = set(args.force_retrain_init_modes)

    rows = []
    model_metrics = {}
    run_subset_slugs = set(args.subset_slugs)
    run_subset_defs = [(idx, item) for idx, item in enumerate(subset_defs) if item[0] in run_subset_slugs]
    init_seed_offsets = {"finetune": 0, "scratch": 10000}
    for init_mode in args.init_modes:
        mode_steps = train_steps_by_init_mode[init_mode]
        for idx, (subset_slug, label, p_thr, s_thr) in run_subset_defs:
            subset_records = train_pools[subset_slug]
            model_slug = f"{init_mode}_{subset_slug}"
            model_label = f"{init_mode} {label if args.no_match_train_size else f'{label} matched'}"
            model = train_model(
                h5_path=h5_path,
                records=subset_records,
                access_index=train_access_index,
                base_ckpt=base_ckpt,
                out_ckpt=out_dir / f"pnsn.v3.{init_mode}.{subset_slug}.pt",
                log_csv=out_dir / f"loss_{init_mode}_{subset_slug}.csv",
                seed=args.seed + init_seed_offsets[init_mode] + idx * 1000,
                steps=mode_steps,
                batch_size=args.train_batch,
                length=args.length,
                padlen=args.padlen,
                lr=args.lr,
                device=device,
                resume=args.resume and init_mode not in force_retrain_init_modes,
                init_mode=init_mode,
                checkpoint_every_steps=args.checkpoint_every_steps,
            )
            if args.train_only:
                continue
            eval_waves, eval_labels, eval_kinds = get_eval_data()
            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[model_slug] = strip_errors(metrics)
            row = {
                "slug": model_slug,
                "subset_slug": subset_slug,
                "init_mode": init_mode,
                "label": model_label,
                "p_threshold_db": p_thr,
                "s_threshold_db": s_thr,
                "train_records": len(subset_records),
                "train_steps": mode_steps,
                "force_retrained": init_mode in force_retrain_init_modes,
                "candidate_train_records": candidate_counts[subset_slug],
                "candidate_phase_counts": candidate_phase_counts[subset_slug],
                "candidate_phase_composition_counts": candidate_phase_composition_counts[subset_slug],
                "matched_phase_counts": matched_phase_counts[subset_slug],
                "matched_phase_composition_counts": matched_phase_composition_counts[subset_slug],
                **metric_row(metrics),
            }
            rows.append(row)
            print(f"metrics {model_label}: {json.dumps(row, ensure_ascii=False)}", flush=True)

    if args.train_only:
        print("train-only run complete; skipped evaluation and summary plotting", flush=True)
        return

    if args.merge_existing_summary and (out_dir / "summary.json").exists():
        existing_summary = json.loads((out_dir / "summary.json").read_text(encoding="utf-8"))
        merged_rows_by_slug = {row["slug"]: row for row in existing_summary.get("rows", [])}
        for row in rows:
            merged_rows_by_slug[row["slug"]] = row
        rows = list(merged_rows_by_slug.values())
        model_metrics = {
            **existing_summary.get("metrics_by_model", {}),
            **model_metrics,
        }

    plot_summary(rows, out_dir / "snr_transfer_phase_balanced_f1_summary.png")
    write_metrics_table(rows, out_dir / "metrics_table.tex")
    summary = {
        "experiment": "phase_balanced_snr_transfer",
        "seed": args.seed,
        "device": str(device),
        "train_steps": args.train_steps,
        "scratch_train_steps": args.scratch_train_steps if args.scratch_train_steps > 0 else args.train_steps,
        "train_steps_by_init_mode": train_steps_by_init_mode,
        "train_batch": args.train_batch,
        "init_modes": args.init_modes,
        "force_retrain_init_modes": sorted(force_retrain_init_modes),
        "checkpoint_every_steps": args.checkpoint_every_steps,
        "subset_slugs": args.subset_slugs,
        "merge_existing_summary": args.merge_existing_summary,
        "filter_mode": args.filter_mode,
        "s_threshold_mode": args.s_threshold_mode,
        "match_mode": args.match_mode,
        "base_ckpt": str(base_ckpt.resolve()),
        "cache_dir": str(cache_dir.resolve()),
        "matched_train_records": matched_train_records,
        "matched_train_composition": matched_train_composition,
        "thresholds": {
            "medium": {"P_db": 5.0, "S_db": s5, "target_P_labels": p5_count, "retained_S_labels": s5_count},
            "high": {"P_db": 10.0, "S_db": s10, "target_P_labels": p10_count, "retained_S_labels": s10_count},
        },
        "snr_summary_db": {"P": summarize_snr(by_group["P"]), "S": summarize_snr(by_group["S"])},
        "candidate_train_records": candidate_counts,
        "candidate_phase_counts": candidate_phase_counts,
        "candidate_phase_composition_counts": candidate_phase_composition_counts,
        "matched_phase_counts": matched_phase_counts,
        "matched_phase_composition_counts": matched_phase_composition_counts,
        "eval_samples": args.eval_samples,
        "eval_samples_single": int(sum(k == "single" for k in get_eval_data()[2])),
        "eval_samples_double": int(sum(k == "double" for k in get_eval_data()[2])),
        "rows": rows,
        "metrics_by_model": model_metrics,
        "phase_balance_table": str((out_dir / "phase_balance_table.csv").resolve()),
        "figure": str((out_dir / "snr_transfer_phase_balanced_f1_summary.png").resolve()),
        "table": str((out_dir / "metrics_table.tex").resolve()),
    }
    (out_dir / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(summary, ensure_ascii=False, indent=2)[:5000], flush=True)


if __name__ == "__main__":
    main()