#!/usr/bin/env python3
"""Generate GRL revision figures from bootstrap and SNR-stratified tables."""

from __future__ import annotations

import argparse
import csv
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


DEFAULT_REVISION_DIR = Path("outputs/grl_revision_20260610")
DEFAULT_FIG_DIR = Path("grl_overleaf/figures/grl_revision")


def read_rows(path: Path) -> list[dict[str, str]]:
    if not path.exists():
        raise FileNotFoundError(f"Missing required table: {path}")
    with path.open(newline="", encoding="utf-8") as f:
        return list(csv.DictReader(f))


def style(ax):
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.grid(axis="y", color="#e6e6e6")
    ax.set_axisbelow(True)


def lookup(rows: list[dict[str, str]], condition: str, metric: str) -> dict[str, float]:
    row = next(r for r in rows if r["condition"] == condition and r["metric"] == metric)
    return {k: float(row[k]) for k in ("estimate", "ci_low", "ci_high")}


def fig2(revision_dir: Path, fig_dir: Path) -> None:
    phase = read_rows(revision_dir / "tables" / "phase_bootstrap_ci.csv")
    disp = read_rows(revision_dir / "tables" / "dispersion_bootstrap_ci.csv")
    fig, axs = plt.subplots(1, 2, figsize=(7.2, 3.1), constrained_layout=True)
    conds = ["full", "snr5", "snr10"]
    labels = ["Full", "SNR>5", "SNR>10"]
    x = np.arange(len(conds))
    ax = axs[0]
    for metric, color, marker in [("P_f1", "#4477AA", "o"), ("S_f1", "#CC6677", "s"), ("mean_f1", "#228833", "^")]:
        vals = [lookup(phase, c, metric) for c in conds]
        y = np.array([v["estimate"] for v in vals])
        yerr = np.array([[v["estimate"] - v["ci_low"] for v in vals], [v["ci_high"] - v["estimate"] for v in vals]])
        ax.errorbar(x, y, yerr=yerr, marker=marker, linewidth=2, capsize=3, label=metric.replace("_", " "), color=color)
    ax.set_xticks(x, labels)
    ax.set_ylabel("F1")
    ax.set_title("Phase picking")
    ax.legend(frameon=False, fontsize=8)
    style(ax)

    ax = axs[1]
    conds_d = ["full", "snr_q1", "snr_q2"]
    labels_d = ["Full", "SNR>3.04", "SNR>6.77"]
    x = np.arange(len(conds_d))
    for metric, color, marker in [("mae", "#4477AA", "o"), ("rmse", "#CC6677", "s")]:
        vals = [lookup(disp, c, metric) for c in conds_d]
        y = np.array([v["estimate"] for v in vals])
        yerr = np.array([[v["estimate"] - v["ci_low"] for v in vals], [v["ci_high"] - v["estimate"] for v in vals]])
        ax.errorbar(x, y, yerr=yerr, marker=marker, linewidth=2, capsize=3, label=metric.upper(), color=color)
    ax.set_xticks(x, labels_d)
    ax.set_ylabel("Velocity error (km/s)")
    ax.set_title("Dispersion")
    ax.legend(frameon=False, fontsize=8)
    style(ax)
    fig_dir.mkdir(parents=True, exist_ok=True)
    for ext in ("pdf", "png"):
        fig.savefig(fig_dir / f"fig2_bootstrap_ci.{ext}", dpi=300, bbox_inches="tight")
    plt.close(fig)


def si_stratified(revision_dir: Path, fig_dir: Path) -> None:
    phase = read_rows(revision_dir / "tables" / "phase_snr_stratified_metrics.csv")
    disp = read_rows(revision_dir / "tables" / "dispersion_snr_stratified_metrics.csv")
    fig, axs = plt.subplots(1, 2, figsize=(7.2, 3.1), constrained_layout=True)
    ax = axs[0]
    for condition, color in [("full", "#4477AA"), ("snr5", "#CC6677"), ("snr10", "#228833")]:
        rows = [r for r in phase if r["condition"] == condition and r["phase"] == "P_S_mean"]
        rows.sort(key=lambda r: r["test_snr_bin"])
        ax.plot(range(len(rows)), [float(r["f1"]) for r in rows], marker="o", label=condition, color=color)
    ax.set_xticks(range(len({r["test_snr_bin"] for r in phase})))
    ax.set_xticklabels(sorted({r["test_snr_bin"] for r in phase}), rotation=25, ha="right", fontsize=7)
    ax.set_ylabel("Mean F1")
    ax.set_title("Phase picking by test SNR")
    ax.legend(frameon=False, fontsize=8)
    style(ax)

    ax = axs[1]
    for condition, color in [("full", "#4477AA"), ("snr_q1", "#CC6677"), ("snr_q2", "#228833")]:
        rows = [r for r in disp if r["condition"] == condition]
        rows.sort(key=lambda r: r["test_snr_bin"])
        ax.plot(range(len(rows)), [float(r["mae"]) for r in rows], marker="o", label=condition, color=color)
    ax.set_xticks(range(len({r["test_snr_bin"] for r in disp})))
    ax.set_xticklabels(sorted({r["test_snr_bin"] for r in disp}), rotation=25, ha="right", fontsize=7)
    ax.set_ylabel("MAE (km/s)")
    ax.set_title("Dispersion by test SNR")
    ax.legend(frameon=False, fontsize=8)
    style(ax)
    fig_dir.mkdir(parents=True, exist_ok=True)
    for ext in ("pdf", "png"):
        fig.savefig(fig_dir / f"figS_snr_stratified.{ext}", dpi=300, bbox_inches="tight")
    plt.close(fig)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--revision-dir", type=Path, default=DEFAULT_REVISION_DIR)
    parser.add_argument("--fig-dir", type=Path, default=DEFAULT_FIG_DIR)
    args = parser.parse_args()
    fig2(args.revision_dir, args.fig_dir)
    si_stratified(args.revision_dir, args.fig_dir)


if __name__ == "__main__":
    main()