#!/usr/bin/env python3
"""Build Figure 4: matched learning performance and continuous pick-recall sweep."""

from __future__ import annotations

import csv
import json
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


OVERLEAF_ROOT = Path(__file__).resolve().parents[1]
PROJECT_ROOT = OVERLEAF_ROOT.parent
OUTPUTS = PROJECT_ROOT / "outputs"
FIG_DIR = OVERLEAF_ROOT / "figures"

PHASE_RECORDS = OUTPUTS / "snr_transfer_phase_balanced_cache" / "records_train_all.json"
PHASE_SNR = OUTPUTS / "snr_transfer_phase_balanced_cache" / "train_phase_snr_db.json"
DISP_SNR = OUTPUTS / "disp_snr_transfer_seed20260609" / "ncf_snr_cache.json"
PHASE_SUMMARY = OUTPUTS / "grl_phase_any_multiseed_seed20260609_20260611" / "phase_multiseed_summary.csv"
PHASE_DIRECT_BASELINE = OUTPUTS / "grl_phase_any_multiseed_seed20260609_20260611" / "phase_direct_baseline_summary.csv"
DISP_SUMMARY = OUTPUTS / "grl_multiseed_seed20260609_20260611" / "dispersion_multiseed_summary.csv"
CONTINUOUS_SWEEP = OUTPUTS / "continuous_phase_recall_snr_conf_sweep" / "continuous_phase_recall_snr_conf_sweep.csv"

OUT_PDF = FIG_DIR / "fig_learning_selection_generalization_summary_v2.pdf"
OUT_PNG = FIG_DIR / "fig_learning_selection_generalization_summary_v2.png"
OUT_DATA = FIG_DIR / "fig_learning_selection_generalization_summary_v2_data.csv"


BLUE = "#1f5a99"
RED = "#b84a4a"
TEAL = "#2a8c8c"
GRAY = "#6d6d6d"
LIGHT_GRAY = "#d8d8d8"
PHASE_TO_GROUP = {"Pg": "P", "Pn": "P", "Sg": "S", "Sn": "S"}
PHASE_THRESHOLDS = {
    "P/S>=5 dB record-any": 5.0,
    "P/S>=10 dB record-any": 10.0,
}


def read_summary(path: Path) -> dict[tuple[str, str], tuple[float, float]]:
    values: dict[tuple[str, str], tuple[float, float]] = {}
    with path.open(newline="", encoding="utf-8") as f:
        for row in csv.DictReader(f):
            values[(row["condition_slug"], row["metric"])] = (
                float(row["mean"]),
                float(row["std"]),
            )
    return values


def binned_fraction(
    distances: list[float],
    snrs: list[float],
    threshold: float | None,
    bins: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    distances_np = np.asarray(distances, dtype=float)
    snrs_np = np.asarray(snrs, dtype=float)
    full_counts, _ = np.histogram(distances_np, bins=bins)
    if threshold is None:
        kept_counts = full_counts.copy()
    else:
        kept_counts, _ = np.histogram(distances_np[snrs_np > threshold], bins=bins)
    with np.errstate(divide="ignore", invalid="ignore"):
        frac = np.where(full_counts > 0, kept_counts / full_counts * 100.0, np.nan)
    centers = (bins[:-1] + bins[1:]) / 2.0
    return centers, frac, full_counts


def binned_mask_fraction(
    distances: list[float],
    keep: list[bool],
    bins: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    distances_np = np.asarray(distances, dtype=float)
    keep_np = np.asarray(keep, dtype=bool)
    full_counts, _ = np.histogram(distances_np, bins=bins)
    kept_counts, _ = np.histogram(distances_np[keep_np], bins=bins)
    with np.errstate(divide="ignore", invalid="ignore"):
        frac = np.where(full_counts > 0, kept_counts / full_counts * 100.0, np.nan)
    centers = (bins[:-1] + bins[1:]) / 2.0
    return centers, frac, full_counts


def phase_snr_key(record: dict, pick: dict, pick_index: int) -> str:
    return (
        f"{record['event']}/{record['station']}/{pick_index}:"
        f"{pick['phase']}:{pick['index']}:{pick['source']}"
    )


def phase_retention() -> dict[str, tuple[np.ndarray, np.ndarray, np.ndarray]]:
    records = json.loads(PHASE_RECORDS.read_text(encoding="utf-8"))["records"]
    snr_by_phase = json.loads(PHASE_SNR.read_text(encoding="utf-8"))
    distances: list[float] = []
    snrs_by_record: list[list[float]] = []
    for record in records:
        record_snrs: list[float] = []
        for pick_index, pick in enumerate(record["phases"]):
            snr = snr_by_phase.get(phase_snr_key(record, pick, pick_index))
            if snr is None:
                continue
            record_snrs.append(float(snr))
        if record_snrs:
            distances.append(float(record["distance_km"]))
            snrs_by_record.append(record_snrs)
    bins = np.arange(0, 451, 50)
    out = {"Full pool": binned_mask_fraction(distances, [True] * len(distances), bins)}
    for label, threshold in PHASE_THRESHOLDS.items():
        keep = [any(np.isfinite(snr) and snr >= threshold for snr in record_snrs) for record_snrs in snrs_by_record]
        out[label] = binned_mask_fraction(distances, keep, bins)
    return {
        key: out[key]
        for key in ["Full pool", "P/S>=5 dB record-any", "P/S>=10 dB record-any"]
    }


def dispersion_retention() -> dict[str, tuple[np.ndarray, np.ndarray, np.ndarray]]:
    cache = json.loads(DISP_SNR.read_text(encoding="utf-8"))
    train_rows = [v for v in cache.values() if v.get("split") == "train"]
    distances = [float(v["distance_km"]) for v in train_rows]
    snrs = [float(v["snr_db"]) for v in train_rows]
    bins = np.arange(75, 576, 50)
    return {
        "Full pool": binned_fraction(distances, snrs, None, bins),
        "SNR>3.04 dB": binned_fraction(distances, snrs, 3.0356278596583564, bins),
        "SNR>6.77 dB": binned_fraction(distances, snrs, 6.774182330903824, bins),
    }


def write_plot_data(
    phase_summary: dict[tuple[str, str], tuple[float, float]],
    direct_baseline: dict[tuple[str, str], tuple[float, float]],
    disp_summary: dict[tuple[str, str], tuple[float, float]],
    continuous_rows: list[dict[str, str]],
) -> None:
    rows: list[dict[str, str | float]] = []
    for condition, label in [
        ("finetune_full", "Fine-tune 2k full"),
        ("finetune_p5_s_bal", "Fine-tune 2k P/S>=5 dB record-any"),
        ("finetune_p10_s_bal", "Fine-tune 2k P/S>=10 dB record-any"),
        ("scratch_full", "Scratch 10k full"),
        ("scratch_p5_s_bal", "Scratch 10k P/S>=5 dB record-any"),
        ("scratch_p10_s_bal", "Scratch 10k P/S>=10 dB record-any"),
    ]:
        for metric in ["P_f1", "S_f1"]:
            mean, std = phase_summary[(condition, metric)]
            rows.append(
                {
                    "panel": "phase_test",
                    "condition": label,
                    "x": condition,
                    "metric": metric,
                    "value": mean,
                    "std": std,
                    "full_bin_count": "",
                }
            )
    for metric in ["P_f1", "S_f1"]:
        mean, std = direct_baseline[("pretrained_direct", metric)]
        rows.append(
            {
                "panel": "phase_test",
                "condition": "Pretrained direct use",
                "x": "direct",
                "metric": metric,
                "value": mean,
                "std": std,
                "full_bin_count": "",
            }
        )
    for condition, label in [
        ("full", "Full distribution"),
        ("snr_q1", "SNR>3.04 dB"),
        ("snr_q2", "SNR>6.77 dB"),
    ]:
        for metric in ["val_mae", "val_rmse"]:
            mean, std = disp_summary[(condition, metric)]
            rows.append(
                {
                    "panel": "dispersion_test",
                    "condition": label,
                    "x": condition,
                    "metric": metric,
                    "value": mean,
                    "std": std,
                    "full_bin_count": "",
                    }
                )
    for row in continuous_rows:
        if row["selection"] == "snr":
            label = "SNR filter"
        elif row["confidence_scope"] == "phase_count_matched":
            label = "Confidence, same phase count"
        else:
            continue
        rows.append(
            {
                "panel": "continuous_phase_recall",
                "condition": label,
                "x": row["threshold_label"],
                "metric": f"{row['phase']}_recall",
                "value": float(row["recall"]),
                "std": "",
                "full_bin_count": row["auto_count_phase"],
            }
        )
    full_phase_counts = {
        row["phase"]: float(row["auto_count_phase"])
        for row in continuous_rows
        if row["selection"] == "snr" and row["threshold_label"] == "Full"
    }
    for row in continuous_rows:
        if row["selection"] != "snr":
            continue
        phase = row["phase"]
        rows.append(
            {
                "panel": "continuous_phase_recall",
                "condition": "Retained picks",
                "x": row["threshold_label"],
                "metric": f"{phase}_retained_pick_percent",
                "value": 100.0 * float(row["auto_count_phase"]) / full_phase_counts[phase],
                "std": "",
                "full_bin_count": row["auto_count_phase"],
            }
        )
    with OUT_DATA.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(
            f,
            fieldnames=["panel", "condition", "x", "metric", "value", "std", "full_bin_count"],
        )
        writer.writeheader()
        writer.writerows(rows)


def style_axes(ax: plt.Axes) -> None:
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.grid(axis="y", color="#ececec", linewidth=0.8)
    ax.tick_params(axis="both", labelsize=8)


def plot_retention(ax: plt.Axes, data: dict[str, tuple[np.ndarray, np.ndarray, np.ndarray]], title: str, xlabel: str) -> None:
    styles = {
        "Full pool": (GRAY, "--", "o"),
        "SNR>5 dB": (BLUE, "-", "o"),
        "SNR>10 dB": (RED, "-", "s"),
        "P/S>=5 dB record-any": (BLUE, "-", "o"),
        "P/S>=10 dB record-any": (RED, "-", "s"),
        "SNR>3.04 dB": (BLUE, "-", "o"),
        "SNR>6.77 dB": (RED, "-", "s"),
    }
    for label, (centers, frac, counts) in data.items():
        mask = counts > 0
        color, linestyle, marker = styles[label]
        ax.plot(
            centers[mask],
            frac[mask],
            color=color,
            linestyle=linestyle,
            marker=marker,
            markersize=3.5,
            linewidth=1.4,
            label=label,
        )
    ax.set_title(title, loc="left", fontsize=10, fontweight="bold")
    ax.set_xlabel(xlabel, fontsize=9)
    ax.set_ylabel("Training candidates retained (%)", fontsize=9)
    ax.set_ylim(-2, 105)
    ax.set_yticks([0, 25, 50, 75, 100])
    style_axes(ax)
    ax.legend(frameon=False, fontsize=7.2, loc="lower left")


def plot_phase_metrics(
    ax: plt.Axes,
    summary: dict[tuple[str, str], tuple[float, float]],
    direct_baseline: dict[tuple[str, str], tuple[float, float]],
) -> None:
    conditions = ["full", "p5_s_bal", "p10_s_bal"]
    labels = ["No filter", "SNR >= 5 dB", "SNR >= 10 dB"]
    x = np.arange(len(conditions), dtype=float)
    phases = [("P", BLUE, -0.045), ("S", RED, 0.045)]
    init_styles = [
        ("finetune", "Fine-tune (2k)", "-", "o", 1.0),
        ("scratch", "Scratch (10k)", "--", "s", 0.86),
    ]
    for phase, color, offset in phases:
        for init_mode, init_label, linestyle, marker, alpha in init_styles:
            means = [summary[(f"{init_mode}_{c}", f"{phase}_f1")][0] for c in conditions]
            stds = [summary[(f"{init_mode}_{c}", f"{phase}_f1")][1] for c in conditions]
            label = f"{phase} {init_label}"
            ax.errorbar(
                x + offset,
                means,
                yerr=stds,
                color=color,
                linestyle=linestyle,
                marker=marker,
                markersize=3.8,
                linewidth=1.25,
                capsize=2.4,
                alpha=alpha,
                label=label,
            )
    for phase, color in [("P", BLUE), ("S", RED)]:
        mean, _std = direct_baseline[("pretrained_direct", f"{phase}_f1")]
        ax.axhline(
            mean,
            color=color,
            linestyle=":",
            linewidth=1.15,
            alpha=0.62,
            label=f"{phase} direct",
        )
    ax.set_title("Phase-picking test performance", loc="left", fontsize=10, fontweight="bold")
    ax.set_ylabel("Phase F1 on unfiltered test windows", fontsize=9)
    ax.set_xticks(x)
    ax.set_xticklabels(labels, fontsize=7.8)
    ax.set_ylim(0.5, 0.77)
    ax.set_yticks([0.5, 0.6, 0.7])
    style_axes(ax)
    ax.legend(frameon=False, fontsize=6.5, loc="lower left", ncol=2, columnspacing=0.75, handlelength=1.6)
    ax.text(2.08, 0.535, "Lower\nF1", fontsize=7.4, color=RED, ha="left", va="bottom")


def plot_dispersion_metrics(ax: plt.Axes, summary: dict[tuple[str, str], tuple[float, float]]) -> None:
    conditions = ["full", "snr_q1", "snr_q2"]
    labels = ["No filter", "SNR > 3.04 dB", "SNR > 6.77 dB"]
    x = np.arange(len(conditions), dtype=float)
    metrics = [("val_mae", "MAE", BLUE, "o"), ("val_rmse", "RMSE", RED, "s")]
    offsets = [-0.09, 0.09]
    for offset, (metric, label, color, marker) in zip(offsets, metrics):
        means = [summary[(c, metric)][0] for c in conditions]
        stds = [summary[(c, metric)][1] for c in conditions]
        ax.errorbar(
            x + offset,
            means,
            yerr=stds,
            color=color,
            marker=marker,
            markersize=4,
            linewidth=1.3,
            capsize=2.5,
            label=label,
        )
    ax.set_title("Dispersion test performance", loc="left", fontsize=10, fontweight="bold")
    ax.set_ylabel("Error on unfiltered test set (km/s)", fontsize=9)
    ax.set_xticks(x)
    ax.set_xticklabels(labels, fontsize=7.8)
    ax.set_ylim(0.043, 0.072)
    style_axes(ax)
    ax.legend(frameon=False, fontsize=7.2, loc="upper left")
    ax.text(2.08, 0.068, "Higher\nerror", fontsize=7.4, color=RED, ha="left", va="center")


def load_continuous_recall() -> list[dict[str, str]]:
    with CONTINUOUS_SWEEP.open(newline="", encoding="utf-8") as f:
        return list(csv.DictReader(f))


def plot_continuous_phase_recall(ax: plt.Axes, rows: list[dict[str, str]], phase: str, show_legend: bool) -> None:
    phase_color = BLUE if phase == "P" else RED
    x_ticks = [-1, 0, 2, 4, 6, 8, 10]
    x_labels = ["No\nfilter", "0", "2", "4", "6", "8", "10"]
    styles = [
        ("snr", "", "SNR filter", "-", "o", 1.0),
        ("confidence", "phase_count_matched", "Confidence, same phase count", "--", "s", 0.72),
    ]
    legend_handles = []
    legend_labels = []
    for selection, scope, label, linestyle, marker, alpha in styles:
        series = sorted(
            [
                row
                for row in rows
                if row["phase"] == phase
                and row["selection"] == selection
                and row["confidence_scope"] == scope
                and (row["threshold_label"] == "Full" or float(row["threshold_db"]).is_integer())
            ],
            key=lambda row: float(row["threshold_db"]),
        )
        (line,) = ax.plot(
            [float(row["threshold_db"]) for row in series],
            [100.0 * float(row["recall"]) for row in series],
            color=phase_color,
            linestyle=linestyle,
            marker=marker,
            markersize=3.7,
            linewidth=1.5,
            alpha=alpha,
            label=label,
        )
        legend_handles.append(line)
        legend_labels.append(label)
    snr_series = sorted(
        [
            row
            for row in rows
            if row["phase"] == phase
            and row["selection"] == "snr"
            and (row["threshold_label"] == "Full" or float(row["threshold_db"]).is_integer())
        ],
        key=lambda row: float(row["threshold_db"]),
    )
    full_count = next(float(row["auto_count_phase"]) for row in snr_series if row["threshold_label"] == "Full")
    retained_percent = [100.0 * float(row["auto_count_phase"]) / full_count for row in snr_series]
    for marker_x in (5.0, 10.0):
        ax.axvline(marker_x, color="#8a8a8a", linewidth=0.75, linestyle=":", zorder=0)
    ax.set_title(f"{phase}-phase continuous recall", loc="left", fontsize=10, fontweight="bold")
    ax.set_xlabel("SNR threshold (dB)", fontsize=9)
    ax.set_ylabel("Recall of covered manual picks (%)", fontsize=9)
    ax.set_xlim(-1.25, 10.2)
    ax.set_xticks(x_ticks)
    ax.set_xticklabels(x_labels, fontsize=8)
    ax.set_ylim(-3, 95)
    style_axes(ax)
    ax_count = ax.twinx()
    (count_line,) = ax_count.plot(
        [float(row["threshold_db"]) for row in snr_series],
        retained_percent,
        color=GRAY,
        linestyle=":",
        marker="^",
        markersize=3.6,
        linewidth=1.45,
        alpha=0.95,
        label="Retained picks (right axis)",
    )
    ax_count.set_ylim(-3, 105)
    ax_count.set_yticks([0, 25, 50, 75, 100])
    ax_count.set_ylabel("Retained picks (%)", fontsize=9, color=GRAY)
    ax_count.tick_params(axis="y", labelsize=8, colors=GRAY)
    ax_count.spines["top"].set_visible(False)
    ax_count.spines["left"].set_visible(False)
    ax_count.spines["right"].set_color(GRAY)
    if show_legend:
        ax.legend(
            legend_handles + [count_line],
            legend_labels + ["Retained picks"],
            frameon=False,
            fontsize=7.0,
            loc="lower left",
        )


def add_panel_label(ax: plt.Axes, label: str) -> None:
    ax.text(-0.13, 1.08, label, transform=ax.transAxes, fontsize=11, fontweight="bold", va="bottom")


def main() -> None:
    phase_summary = read_summary(PHASE_SUMMARY)
    direct_baseline = read_summary(PHASE_DIRECT_BASELINE)
    disp_summary = read_summary(DISP_SUMMARY)
    continuous_rows = load_continuous_recall()
    write_plot_data(phase_summary, direct_baseline, disp_summary, continuous_rows)

    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 8.5,
            "axes.linewidth": 0.8,
            "pdf.fonttype": 42,
            "ps.fonttype": 42,
        }
    )
    fig, axes = plt.subplots(2, 2, figsize=(7.2, 5.0), constrained_layout=True)

    plot_phase_metrics(axes[0, 0], phase_summary, direct_baseline)
    plot_dispersion_metrics(axes[0, 1], disp_summary)
    plot_continuous_phase_recall(axes[1, 0], continuous_rows, "P", show_legend=True)
    plot_continuous_phase_recall(axes[1, 1], continuous_rows, "S", show_legend=False)

    for ax, label in zip(axes.flat, ["a", "b", "c", "d"]):
        add_panel_label(ax, label)

    fig.savefig(OUT_PDF, bbox_inches="tight")
    fig.savefig(OUT_PNG, dpi=300, bbox_inches="tight")
    print(f"Wrote {OUT_PDF}")
    print(f"Wrote {OUT_PNG}")
    print(f"Wrote {OUT_DATA}")


if __name__ == "__main__":
    main()