#!/usr/bin/env python3
"""Regenerate all main-manuscript figures from archived plotted data.

This script is intentionally data-light. It reads the CSV/JSON plotted-data
exports stored in ``results/manuscript_figures`` and writes regenerated PDF and
PNG figures to ``reproduced_figures``. Full recomputation from waveform and pick
data is handled by the training, evaluation, and association scripts elsewhere
in this archive.
"""

from __future__ import annotations

import argparse
import csv
import json
from collections import Counter
from pathlib import Path

import matplotlib

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


BLUE = "#2F6DB2"
RED = "#C44E52"
ORANGE = "#F58518"
GRAY = "#6D6D6D"
LIGHT = "#E8E8E8"
DARK = "#222222"


def read_csv(path: Path) -> list[dict[str, str]]:
    with path.open(newline="", encoding="utf-8") as handle:
        return list(csv.DictReader(handle))


def read_json(path: Path) -> dict:
    with path.open(encoding="utf-8") as handle:
        return json.load(handle)


def style_axis(ax: plt.Axes, grid_axis: str = "y") -> None:
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    if grid_axis:
        ax.grid(axis=grid_axis, color=LIGHT, linewidth=0.7, alpha=0.9)
        ax.set_axisbelow(True)
    ax.tick_params(labelsize=8, width=0.7, length=3)


def panel_label(ax: plt.Axes, label: str) -> None:
    ax.text(
        -0.12,
        1.06,
        label,
        transform=ax.transAxes,
        fontsize=12,
        fontweight="bold",
        va="bottom",
        color=DARK,
    )


def numeric_threshold(value: str) -> float:
    return -1.0 if value == "Full" else float(value)


def figure_observability(src: Path, out_dir: Path) -> None:
    rows = read_csv(src / "fig_observability_real_data_v1_data.csv")
    summary = read_json(src / "fig_observability_real_data_v1_summary.json")
    phase = [r for r in rows if r["panel"] == "phase_scatter"]
    distance = [r for r in rows if r["panel"] == "distance_coverage"]

    plt.rcParams.update({"font.family": "DejaVu Sans", "pdf.fonttype": 42, "ps.fonttype": 42})
    fig, axes = plt.subplots(1, 3, figsize=(7.25, 3.05), constrained_layout=True, dpi=300)

    ax = axes[0]
    for metric, color, label in [("P_pick_snr_db", RED, "P picks"), ("S_pick_snr_db", BLUE, "S picks")]:
        pts = [r for r in phase if r["metric"] == metric]
        ax.scatter(
            [float(r["x"]) for r in pts],
            [min(float(r["y"]), 22.0) for r in pts],
            s=13,
            color=color,
            alpha=0.68,
            linewidth=0,
            label=label,
        )
    ax.axhline(5.0, color="#8B0000", linestyle="--", linewidth=1.0)
    ax.set_xscale("log")
    ax.set_xlim(4, 1000)
    ax.set_ylim(-2.5, 22.5)
    ax.set_xlabel("Epicentral distance (km)")
    ax.set_ylabel("Pick SNR (dB)")
    ax.set_title("Unfiltered phase picks", loc="left", fontweight="bold", fontsize=9)
    ax.legend(frameon=False, fontsize=7, loc="upper right")
    style_axis(ax, "both")
    panel_label(ax, "a")

    ax = axes[1]
    bins = summary["dispersion_period_bins"]
    x = np.arange(len(bins))
    med = np.array([b["median"] for b in bins], dtype=float)
    q25 = np.array([b["q25"] for b in bins], dtype=float)
    q75 = np.array([b["q75"] for b in bins], dtype=float)
    q05 = np.array([b["q05"] for b in bins], dtype=float)
    q95 = np.array([b["q95"] for b in bins], dtype=float)
    ax.vlines(x, q05, q95, color=GRAY, linewidth=1.2, alpha=0.8)
    ax.bar(x, q75 - q25, bottom=q25, color="#9FC6DF", edgecolor="#3C6682", linewidth=0.6, alpha=0.75)
    ax.scatter(x, med, color="#1F4E79", s=18, zorder=3, label="median")
    ax.axhline(5.0, color="#8B0000", linestyle="--", linewidth=1.0)
    ax.set_xticks(x, [f"{int(b['period_left'])}-{int(b['period_right'])}" for b in bins])
    ax.set_xlabel("Period bin (s)")
    ax.set_ylabel("Dispersion SNR (dB)")
    ax.set_title("Dispersion SNR by period", loc="left", fontweight="bold", fontsize=9)
    style_axis(ax)
    panel_label(ax, "b")

    ax = axes[2]
    labels = [
        ("unfiltered", "Unfiltered", GRAY, "--"),
        ("SNR >= 5 dB", "SNR >= 5 dB", BLUE, "-"),
        ("SNR >= 10 dB", "SNR >= 10 dB", RED, "-"),
    ]
    for condition, label, color, linestyle in labels:
        pts = [r for r in distance if r["condition"] == condition and r["metric"] in {"fraction", "retained_fraction"}]
        pts = sorted(pts, key=lambda r: float(r["x"]))
        if not pts:
            continue
        centers = [(float(r["x"]) + float(r["y"])) / 2.0 for r in pts]
        vals = [100.0 * float(r["value"]) for r in pts]
        ax.plot(centers, vals, color=color, linestyle=linestyle, marker="o", markersize=3.2, linewidth=1.4, label=label)
    ax.set_xlim(0, 450)
    ax.set_ylim(-3, 105)
    ax.set_xlabel("Epicentral distance (km)")
    ax.set_ylabel("Records retained (%)")
    ax.set_title("Distance coverage changes", loc="left", fontweight="bold", fontsize=9)
    ax.legend(frameon=False, fontsize=7, loc="upper right")
    style_axis(ax)
    panel_label(ax, "c")

    fig.suptitle("Hard SNR thresholds reshape seismic observability", fontsize=11, fontweight="bold")
    for suffix in ["pdf", "png"]:
        fig.savefig(out_dir / f"fig1_observability_regenerated.{suffix}", bbox_inches="tight", dpi=300)
    plt.close(fig)


def figure_learning(src: Path, out_dir: Path) -> None:
    rows = read_csv(src / "fig_learning_selection_generalization_summary_v2_data.csv")
    plt.rcParams.update({"font.family": "DejaVu Sans", "pdf.fonttype": 42, "ps.fonttype": 42})
    fig, axes = plt.subplots(2, 2, figsize=(7.25, 4.9), constrained_layout=True, dpi=300)

    phase_rows = [r for r in rows if r["panel"] == "phase_test"]
    phase_order = ["finetune_full", "finetune_p5_s_bal", "finetune_p10_s_bal", "direct", "scratch_full", "scratch_p5_s_bal", "scratch_p10_s_bal"]
    labels = ["FT full", "FT 5 dB", "FT 10 dB", "Direct", "Scratch full", "Scratch 5 dB", "Scratch 10 dB"]
    x = np.arange(len(phase_order), dtype=float)
    ax = axes[0, 0]
    for offset, metric, color, label in [(-0.17, "P_f1", BLUE, "P F1"), (0.17, "S_f1", RED, "S F1")]:
        vals, errs = [], []
        for key in phase_order:
            match = [r for r in phase_rows if r["x"] == key and r["metric"] == metric]
            vals.append(float(match[0]["value"]) if match else np.nan)
            errs.append(float(match[0]["std"]) if match and match[0]["std"] else 0.0)
        ax.bar(x + offset, vals, width=0.32, yerr=errs, capsize=2, color=color, alpha=0.82, label=label)
    ax.set_xticks(x, labels, rotation=35, ha="right")
    ax.set_ylim(0.56, 0.78)
    ax.set_ylabel("Unfiltered-test F1")
    ax.set_title("Phase-picking generalization", loc="left", fontweight="bold", fontsize=9)
    ax.legend(frameon=False, fontsize=7, loc="lower left")
    style_axis(ax)
    panel_label(ax, "a")

    disp_rows = [r for r in rows if r["panel"] == "dispersion_test"]
    disp_order = ["full", "snr_q1", "snr_q2"]
    disp_labels = ["Full", "SNR>3.04", "SNR>6.77"]
    ax = axes[0, 1]
    x = np.arange(len(disp_order), dtype=float)
    for offset, metric, color, label in [(-0.08, "val_mae", BLUE, "MAE"), (0.08, "val_rmse", RED, "RMSE")]:
        vals, errs = [], []
        for key in disp_order:
            match = [r for r in disp_rows if r["x"] == key and r["metric"] == metric]
            vals.append(float(match[0]["value"]))
            errs.append(float(match[0]["std"]))
        ax.errorbar(x + offset, vals, yerr=errs, color=color, marker="o", linewidth=1.4, capsize=2.5, label=label)
    ax.set_xticks(x, disp_labels)
    ax.set_ylim(0.043, 0.072)
    ax.set_ylabel("Error (km/s)")
    ax.set_title("Dispersion generalization", loc="left", fontweight="bold", fontsize=9)
    ax.legend(frameon=False, fontsize=7, loc="upper left")
    style_axis(ax)
    panel_label(ax, "b")

    recall_rows = [r for r in rows if r["panel"] == "continuous_phase_recall"]
    for ax, phase, label in [(axes[1, 0], "P", "c"), (axes[1, 1], "S", "d")]:
        color = BLUE if phase == "P" else RED
        for condition, linestyle, marker in [("SNR filter", "-", "o"), ("Confidence, same phase count", "--", "s")]:
            pts = [
                r
                for r in recall_rows
                if r["condition"] == condition and r["metric"] == f"{phase}_recall"
            ]
            pts = sorted(pts, key=lambda r: numeric_threshold(r["x"]))
            ax.plot(
                [numeric_threshold(r["x"]) for r in pts],
                [100.0 * float(r["value"]) for r in pts],
                color=color,
                linestyle=linestyle,
                marker=marker,
                markersize=3.5,
                linewidth=1.5,
                alpha=0.95 if condition == "SNR filter" else 0.68,
                label=condition,
            )
        retained = sorted(
            [r for r in recall_rows if r["condition"] == "Retained picks" and r["metric"] == f"{phase}_retained_pick_percent"],
            key=lambda r: numeric_threshold(r["x"]),
        )
        ax2 = ax.twinx()
        ax2.plot([numeric_threshold(r["x"]) for r in retained], [float(r["value"]) for r in retained], color=GRAY, linestyle=":", marker="^", markersize=3.2)
        ax2.set_ylim(-3, 105)
        ax2.set_ylabel("Retained picks (%)", color=GRAY)
        ax2.tick_params(axis="y", labelsize=8, colors=GRAY)
        ax2.spines["top"].set_visible(False)
        ax.set_xlim(-1.25, 10.25)
        ax.set_ylim(-3, 95)
        ax.set_xlabel("SNR threshold (dB)")
        ax.set_ylabel("Coverage-filtered label recall (%)")
        ax.set_title(f"{phase}-phase continuous recall", loc="left", fontweight="bold", fontsize=9)
        if phase == "P":
            ax.legend(frameon=False, fontsize=7, loc="lower left")
        style_axis(ax)
        panel_label(ax, label)

    fig.suptitle("SNR filtering effects from training to deployment", fontsize=11, fontweight="bold")
    for suffix in ["pdf", "png"]:
        fig.savefig(out_dir / f"fig2_learning_and_deployment_regenerated.{suffix}", bbox_inches="tight", dpi=300)
    plt.close(fig)


def figure_event_geometry(src: Path, out_dir: Path) -> None:
    rows = read_csv(src / "fig_event_geometry_distribution_polished_data.csv")
    summary = read_json(src / "fig_event_geometry_distribution_polished_summary.json")["conditions"]
    plt.rcParams.update({"font.family": "DejaVu Sans", "pdf.fonttype": 42, "ps.fonttype": 42})
    fig, axes = plt.subplots(2, 2, figsize=(7.25, 4.75), constrained_layout=True, dpi=300)

    panels = [
        (axes[0, 0], "moderate", "S", "a"),
        (axes[0, 1], "strict", "S", "b"),
        (axes[1, 0], "moderate", "S_stations", "c"),
        (axes[1, 1], "strict", "S_stations", "d"),
    ]
    class_labels = {"snr_only": "SNR-only events", "confidence_only": "Confidence-only events"}
    class_colors = {"snr_only": ORANGE, "confidence_only": BLUE}
    for ax, condition, metric, label in panels:
        cond = summary[condition]
        subset = [r for r in rows if r["condition_slug"] == condition and r["metric"] == metric]
        all_values = [int(r["difference"]) for r in subset]
        xmin = min(-8, min(all_values) if all_values else -1)
        xmax = max(16, max(all_values) if all_values else 1)
        xs = np.arange(xmin, xmax + 1)
        width = 0.38
        for class_slug, offset in [("snr_only", -width / 2), ("confidence_only", width / 2)]:
            vals = [int(r["difference"]) for r in subset if r["class_slug"] == class_slug]
            counts = Counter(vals)
            denom = len(vals) or 1
            heights = np.array([counts.get(int(x), 0) / denom for x in xs])
            ax.bar(xs + offset, heights, width=width, color=class_colors[class_slug], edgecolor="black", linewidth=0.25, alpha=0.86, label=class_labels[class_slug])
        ax.axvline(0, color=DARK, linewidth=0.8)
        if metric == "S":
            title = (
                f"{cond['threshold']}, matched S arrivals\n"
                f"P/R: SNR {cond['snr_event_metrics']['precision']:.3f}/{cond['snr_event_metrics']['recall']:.3f}; "
                f"prob. {cond['confidence_event_metrics']['precision']:.3f}/{cond['confidence_event_metrics']['recall']:.3f}"
            )
        else:
            title = f"{cond['threshold']}, S-contributing stations"
        ax.set_title(title, fontsize=8)
        ax.set_xlabel("Recovered stream minus missed stream")
        ax.set_ylabel("Fraction of discordant events")
        ax.set_xlim(xmin - 0.8, xmax + 0.8)
        style_axis(ax)
        panel_label(ax, label)
        if label in {"a", "b"}:
            ax.legend(frameon=False, fontsize=7, loc="upper left")

    fig.suptitle("Phase-balanced event-level retained support", fontsize=11, fontweight="bold")
    for suffix in ["pdf", "png"]:
        fig.savefig(out_dir / f"fig3_phase_balanced_event_geometry_regenerated.{suffix}", bbox_inches="tight", dpi=300)
    plt.close(fig)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    archive_root = Path(__file__).resolve().parents[2]
    parser.add_argument("--source-dir", type=Path, default=archive_root / "results" / "manuscript_figures")
    parser.add_argument("--out-dir", type=Path, default=archive_root / "reproduced_figures")
    args = parser.parse_args()
    args.out_dir.mkdir(parents=True, exist_ok=True)

    figure_observability(args.source_dir, args.out_dir)
    figure_learning(args.source_dir, args.out_dir)
    figure_event_geometry(args.source_dir, args.out_dir)
    print(f"Wrote regenerated figures to {args.out_dir}")


if __name__ == "__main__":
    main()