from __future__ import annotations

import csv
import json
from pathlib import Path

import matplotlib.pyplot as plt
from matplotlib.ticker import PercentFormatter
import numpy as np


ROOT = Path(__file__).resolve().parents[1]
PAPER = ROOT / "grl_overleaf"
FIGDIR = PAPER / "figures"
SUMMARY = (
    ROOT
    / "odata"
    / "snr60_pnsn_v3_5120_20190706_20211113_reals_4_2_3_1_1p0_0p1_1p0"
    / "final_snr_confidence_recall_summary.json"
)
MULTISEED = ROOT / "outputs" / "grl_multiseed_seed20260609_20260611" / "multiseed_summary.json"

BLUE = "#4477AA"
ORANGE = "#CC6677"
GREEN = "#228833"
GRAY = "#666666"


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


def metric_lookup(rows: list[dict[str, str]], metric_key: str) -> dict[str, float]:
    return {r["condition_slug"]: float(r["value"]) for r in rows if r["metric_key"] == metric_key}


def multiseed_lookup(data: dict, task: str, metric: str) -> dict[str, dict[str, object]]:
    out: dict[str, dict[str, object]] = {}
    for row in data["rows"]:
        if row["task"] == task and row["metric"] == metric:
            out[row["condition_slug"]] = row
    return out


def values(row: dict[str, object]) -> np.ndarray:
    return np.asarray(row["values"], dtype=float)


def rel_stats(num: dict[str, object], den: dict[str, object], sign: float = 1.0) -> tuple[float, float]:
    change = sign * (values(num) - values(den)) / values(den)
    return float(np.mean(change)), float(np.std(change, ddof=1))


def style_axes(ax: plt.Axes) -> None:
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.grid(axis="y", color="#E5E5E5", linewidth=0.8)
    ax.set_axisbelow(True)


def panel_label(ax: plt.Axes, text: str) -> None:
    ax.text(
        -0.12,
        1.08,
        text,
        transform=ax.transAxes,
        fontsize=11,
        fontweight="bold",
        va="top",
    )


def make_combined_training_figure() -> None:
    if MULTISEED.exists():
        data = json.loads(MULTISEED.read_text(encoding="utf-8"))
        phase_metrics = {
            key: multiseed_lookup(data, "phase", key)
            for key in ["P_f1", "S_f1", "mean_f1", "P_precision", "S_precision", "P_recall", "S_recall"]
        }
        disp_metrics = {
            key: multiseed_lookup(data, "dispersion", key)
            for key in ["val_mae", "val_rmse"]
        }
    else:
        phase = read_metric_csv(FIGDIR / "fig2_phase_picking_optimized_data.csv")
        disp = read_metric_csv(FIGDIR / "fig3_dispersion_optimized_data.csv")
        phase_metrics = {}
        for key in ["P_f1", "S_f1", "mean_f1", "P_precision", "S_precision", "P_recall", "S_recall"]:
            phase_metrics[key] = {
                cond: {"mean": value, "std": 0.0, "values": [value]}
                for cond, value in metric_lookup(phase, key).items()
            }
        disp_metrics = {}
        for key in ["val_mae", "val_rmse"]:
            disp_metrics[key] = {
                cond: {"mean": value, "std": 0.0, "values": [value]}
                for cond, value in metric_lookup(disp, key).items()
            }

    phase_order = ["full", "snr5", "snr10"]
    phase_labels = ["Full", "SNR>5", "SNR>10"]
    disp_order = ["full", "snr_q1", "snr_q2"]
    disp_labels = ["Full", "SNR>3.04", "SNR>6.77"]

    fig, axs = plt.subplots(2, 2, figsize=(7.2, 6.0), constrained_layout=True)

    ax = axs[0, 0]
    x = np.arange(len(phase_order))
    for key, marker, label, color in [
        ("P_f1", "o", "P", BLUE),
        ("S_f1", "s", "S", ORANGE),
        ("mean_f1", "^", "Mean", GREEN),
    ]:
        rows = phase_metrics[key]
        ax.errorbar(
            x,
            [rows[k]["mean"] for k in phase_order],
            yerr=[rows[k]["std"] for k in phase_order],
            marker=marker,
            linewidth=2,
            capsize=3,
            label=label,
            color=color,
        )
    ax.set_xticks(x, phase_labels)
    ax.set_ylim(0.62, 0.82)
    ax.set_ylabel("F1")
    ax.set_title("Phase-picking F1")
    ax.legend(frameon=False, ncols=3, fontsize=8)
    style_axes(ax)
    panel_label(ax, "A")

    ax = axs[0, 1]
    for key, marker, linestyle, label, color in [
        ("P_precision", "o", "-", "P prec.", BLUE),
        ("P_recall", "o", "--", "P rec.", "#88CCEE"),
        ("S_precision", "s", "-", "S prec.", ORANGE),
        ("S_recall", "s", "--", "S rec.", "#DDCC77"),
    ]:
        rows = phase_metrics[key]
        ax.errorbar(
            x,
            [rows[k]["mean"] for k in phase_order],
            yerr=[rows[k]["std"] for k in phase_order],
            marker=marker,
            linewidth=2,
            linestyle=linestyle,
            capsize=3,
            label=label,
            color=color,
        )
    ax.set_xticks(x, phase_labels)
    ax.set_ylim(0.58, 0.84)
    ax.set_ylabel("Score")
    ax.set_title("Phase-picking precision and recall")
    ax.legend(frameon=False, ncols=2, fontsize=7.5)
    style_axes(ax)
    panel_label(ax, "B")

    ax = axs[1, 0]
    x2 = np.arange(len(disp_order))
    for key, marker, label, color in [
        ("val_mae", "o", "MAE", BLUE),
        ("val_rmse", "s", "RMSE", ORANGE),
    ]:
        rows = disp_metrics[key]
        ax.errorbar(
            x2,
            [rows[k]["mean"] for k in disp_order],
            yerr=[rows[k]["std"] for k in disp_order],
            marker=marker,
            linewidth=2,
            capsize=3,
            label=label,
            color=color,
        )
    ax.set_xticks(x2, disp_labels)
    ax.set_ylim(0.04, 0.073)
    ax.set_ylabel("Velocity error (km/s)")
    ax.set_title("Dispersion estimation")
    ax.legend(frameon=False, ncols=2, fontsize=8)
    style_axes(ax)
    panel_label(ax, "C")

    ax = axs[1, 1]
    stage_labels = ["Full", "Moderate", "Strict"]
    x3 = np.arange(len(stage_labels))
    mean_f1_change = [(0.0, 0.0)] + [
        rel_stats(phase_metrics["mean_f1"][cond], phase_metrics["mean_f1"]["full"])
        for cond in ["snr5", "snr10"]
    ]
    mae_change = [(0.0, 0.0)] + [
        rel_stats(disp_metrics["val_mae"][cond], disp_metrics["val_mae"]["full"])
        for cond in ["snr_q1", "snr_q2"]
    ]
    rmse_change = [(0.0, 0.0)] + [
        rel_stats(disp_metrics["val_rmse"][cond], disp_metrics["val_rmse"]["full"])
        for cond in ["snr_q1", "snr_q2"]
    ]
    ax.axhline(0, color="#333333", linewidth=0.8)
    for series, marker, label, color in [
        (mean_f1_change, "^", "Mean F1", GREEN),
        (mae_change, "o", "MAE", BLUE),
        (rmse_change, "s", "RMSE", ORANGE),
    ]:
        ax.errorbar(
            x3,
            [v[0] for v in series],
            yerr=[v[1] for v in series],
            marker=marker,
            linewidth=2,
            capsize=3,
            label=label,
            color=color,
        )
    ax.set_xticks(x3, stage_labels)
    ax.yaxis.set_major_formatter(PercentFormatter(1.0))
    ax.set_ylabel("Relative change vs full")
    ax.set_title("Trend relative to full training")
    ax.legend(frameon=False, fontsize=8)
    style_axes(ax)
    panel_label(ax, "D")

    for ax in axs.ravel():
        ax.tick_params(labelsize=8)
        ax.title.set_fontsize(10)

    for ext in ["pdf", "png"]:
        fig.savefig(FIGDIR / f"fig2_training_combined.{ext}", dpi=300, bbox_inches="tight")
    plt.close(fig)


def make_association_figure() -> None:
    data = json.loads(SUMMARY.read_text(encoding="utf-8"))
    pick_eval = data["pick_eval"]
    event_eval = data["association"]["event_eval"]

    conds = ["snr", "confidence"]
    labels = ["SNR", "Confidence"]
    colors = [BLUE, ORANGE]
    hatches = ["", "///"]

    fig, axs = plt.subplots(2, 2, figsize=(7.2, 6.0), constrained_layout=True)

    for subset, ax, label, letter in [
        ("all", axs[0, 0], "P/S observability: manual + automatic labels", "A"),
        ("manual", axs[0, 1], "P/S observability: manual labels", "B"),
    ]:
        phases = ["P", "S", "P_S_combined"]
        phase_labels = ["P", "S", "P+S"]
        xx = np.arange(len(phases))
        width = 0.36
        for i, cond in enumerate(conds):
            vals = [
                pick_eval[cond]["recall_covered_subsets"][subset][ph]["recall_covered"]
                for ph in phases
            ]
            ax.bar(
                xx + (i - 0.5) * width,
                vals,
                width,
                label=labels[i],
                color=colors[i],
                hatch=hatches[i],
                edgecolor="#333333",
                linewidth=0.4,
            )
            for xpos, val in zip(xx + (i - 0.5) * width, vals):
                ax.text(xpos, val + 0.018, f"{val:.2f}", ha="center", va="bottom", fontsize=6.8)
        ax.set_xticks(xx, phase_labels)
        ax.set_ylim(0, 0.95)
        ax.set_ylabel("Recall covered")
        ax.set_title(label)
        ax.legend(frameon=False, fontsize=8)
        style_axes(ax)
        panel_label(ax, letter)

    ax = axs[1, 0]
    event_recall = [event_eval[c]["metrics"]["recall"] for c in conds]
    x = np.arange(len(conds))
    bars = ax.bar(x, event_recall, color=colors, edgecolor="#333333", linewidth=0.4)
    for bar, hatch in zip(bars, hatches):
        bar.set_hatch(hatch)
    ax.set_xticks(x, labels)
    ax.set_ylim(0, 0.75)
    ax.set_ylabel("Earthquake recall")
    ax.set_title("Association observability")
    for bar, cond in zip(bars, conds):
        tp = event_eval[cond]["counts"]["true_positive_events"]
        ref = event_eval[cond]["counts"]["reference_events"]
        ax.text(
            bar.get_x() + bar.get_width() / 2,
            bar.get_height() + 0.02,
            f"{tp}/{ref}",
            ha="center",
            va="bottom",
            fontsize=8,
        )
    style_axes(ax)
    panel_label(ax, "C")

    ax = axs[1, 1]
    tp = [event_eval[c]["counts"]["true_positive_events"] for c in conds]
    fn = [event_eval[c]["counts"]["false_negative_events"] for c in conds]
    ax.bar(x, tp, color=GREEN, label="Recovered", edgecolor="#333333", linewidth=0.4)
    ax.bar(x, fn, bottom=tp, color=GRAY, hatch="///", label="Missed", edgecolor="#333333", linewidth=0.4)
    ax.set_xticks(x, labels)
    ax.set_ylim(0, 3000)
    ax.set_ylabel("Catalog events")
    ax.set_title("Recovered vs missed catalog events")
    ax.legend(frameon=False, fontsize=8)
    style_axes(ax)
    panel_label(ax, "D")

    for ax in axs.ravel():
        ax.tick_params(labelsize=8)
        ax.title.set_fontsize(10)

    for ext in ["pdf", "png"]:
        fig.savefig(FIGDIR / f"fig3_continuous_association.{ext}", dpi=300, bbox_inches="tight")
    plt.close(fig)

    rows = []
    for subset in ["all", "manual"]:
        for cond in conds:
            for phase in ["P", "S", "P_S_combined"]:
                rec = pick_eval[cond]["recall_covered_subsets"][subset][phase]["recall_covered"]
                rows.append(
                    {
                        "panel": "phase_recall",
                        "subset": subset,
                        "condition": cond,
                        "phase": phase,
                        "value": rec,
                    }
                )
    for cond in conds:
        rows.append(
            {
                "panel": "event_recall",
                "subset": "catalog",
                "condition": cond,
                "phase": "event",
                "value": event_eval[cond]["metrics"]["recall"],
            }
        )
    with (FIGDIR / "fig3_continuous_association_data.csv").open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(
            f,
            fieldnames=["panel", "subset", "condition", "phase", "value"],
            lineterminator="\n",
        )
        writer.writeheader()
        writer.writerows(rows)


def main() -> None:
    make_combined_training_figure()
    make_association_figure()


if __name__ == "__main__":
    main()