#!/usr/bin/env python3
"""Generate optimized GRL Figures 2 and 3 from matched SNR experiments.

The current experiment directories contain aggregate summaries but not
per-window phase predictions or per-sample dispersion errors. This script
therefore plots verified point estimates and paired differences, while marking
bootstrap confidence intervals as unavailable instead of fabricating them.
"""

from __future__ import annotations

import csv
import json
import math
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterable

import matplotlib.pyplot as plt
import numpy as np


ROOT = Path(__file__).resolve().parents[1]
PHASE_DIR = ROOT / "outputs" / "snr_transfer_seed20260609"
DISP_DIR = ROOT / "outputs" / "disp_snr_transfer_seed20260609"
FIG_DIR = ROOT / "grl_overleaf" / "figures"
REPORT = ROOT / "outputs" / "figure_metric_discrepancy_report.md"

CONDITION_LABELS = {
    "full": "Full distribution",
    "snr5": "Moderate SNR filter",
    "snr10": "Strict SNR filter",
    "snr_q1": "Moderate SNR filter",
    "snr_q2": "Strict SNR filter",
}

PHASE_ORDER = ["full", "snr5", "snr10"]
DISP_ORDER = ["full", "snr_q1", "snr_q2"]

PALETTE = {
    "P F1": "#0072B2",
    "S F1": "#D55E00",
    "Mean F1": "#009E73",
    "P precision": "#0072B2",
    "S precision": "#D55E00",
    "P recall": "#56B4E9",
    "S recall": "#E69F00",
    "MAE": "#0072B2",
    "RMSE": "#D55E00",
}

CI_STATUS = "not_available_missing_per_sample_outputs"
CI_NOTE = "95% CI unavailable from aggregate-only outputs"


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


def write_csv(path: Path, rows: Iterable[dict]) -> None:
    rows = list(rows)
    path.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = sorted({key for row in rows for key in row.keys()})
    with path.open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


def row_map(summary: dict) -> dict[str, dict]:
    return {row["slug"]: row for row in summary["rows"]}


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


def add_panel_label(ax, label: str) -> None:
    ax.text(
        0.01,
        0.94,
        label,
        transform=ax.transAxes,
        fontsize=11,
        fontweight="bold",
        va="top",
        ha="left",
        bbox={"facecolor": "white", "edgecolor": "none", "alpha": 0.85, "pad": 1.5},
    )


def annotate_ci_unavailable(ax) -> None:
    ax.text(
        0.99,
        0.04,
        CI_NOTE,
        transform=ax.transAxes,
        ha="right",
        va="bottom",
        fontsize=7,
        color="#666666",
    )


def plot_line_points(ax, x: np.ndarray, values: list[float], label: str, marker: str = "o") -> None:
    ax.plot(
        x,
        values,
        marker=marker,
        linewidth=1.6,
        markersize=5.2,
        label=label,
        color=PALETTE[label],
    )


def per_sample_candidates(directory: Path) -> list[str]:
    patterns = ("pred", "prediction", "per_sample", "per-window", "per_window", "error", "sample_metrics")
    data_suffixes = {".csv", ".json", ".npz", ".npy", ".parquet", ".h5", ".hdf5"}
    out = []
    for path in directory.glob("*"):
        if path.name.startswith("._"):
            continue
        if path.suffix.lower() not in data_suffixes:
            continue
        lower = path.name.lower()
        if any(pattern in lower for pattern in patterns):
            out.append(str(path.relative_to(ROOT)))
    return sorted(out)


def make_phase_data(phase: dict) -> list[dict]:
    rows = row_map(phase)
    baseline = rows["full"]
    out = []
    for slug in PHASE_ORDER:
        row = rows[slug]
        for metric_key, metric_label in [
            ("P_f1", "P F1"),
            ("S_f1", "S F1"),
            ("mean_f1", "Mean F1"),
            ("P_precision", "P precision"),
            ("S_precision", "S precision"),
            ("P_recall", "P recall"),
            ("S_recall", "S recall"),
        ]:
            value = float(row[metric_key])
            base = float(baseline[metric_key])
            out.append(
                {
                    "figure": "Figure 2",
                    "condition": CONDITION_LABELS[slug],
                    "condition_slug": slug,
                    "condition_index": PHASE_ORDER.index(slug),
                    "metric": metric_label,
                    "metric_key": metric_key,
                    "value": value,
                    "delta_vs_full": value - base,
                    "ci_low": "",
                    "ci_high": "",
                    "delta_ci_low": "",
                    "delta_ci_high": "",
                    "ci_status": CI_STATUS,
                    "ci_method": "paired bootstrap over 10000 windows requires per-window prediction counts",
                    "train_records": row["train_records"],
                    "test_samples": phase["eval_samples"],
                    "test_set": "unfiltered CREDIT-X1 test windows",
                }
            )
    return out


def make_disp_data(disp: dict) -> list[dict]:
    rows = row_map(disp)
    baseline = rows["full"]
    out = []
    for slug in DISP_ORDER:
        row = rows[slug]
        for metric_key, metric_label in [("val_mae", "MAE"), ("val_rmse", "RMSE")]:
            value = float(row[metric_key])
            base = float(baseline[metric_key])
            delta = value - base
            pct = 100.0 * delta / base if base else math.nan
            out.append(
                {
                    "figure": "Figure 3",
                    "condition": CONDITION_LABELS[slug],
                    "condition_slug": slug,
                    "condition_index": DISP_ORDER.index(slug),
                    "metric": metric_label,
                    "metric_key": metric_key,
                    "value": value,
                    "delta_vs_full": delta,
                    "percent_change_vs_full": pct,
                    "ci_low": "",
                    "ci_high": "",
                    "delta_ci_low": "",
                    "delta_ci_high": "",
                    "ci_status": CI_STATUS,
                    "ci_method": "paired bootstrap over 8292 samples requires per-sample prediction errors",
                    "train_records": row["train_records"],
                    "test_samples": disp["test_records"],
                    "test_set": "unfiltered NCF test samples",
                }
            )
    return out


def figure2(phase: dict, data: list[dict]) -> None:
    x = np.arange(len(PHASE_ORDER))
    labels = [CONDITION_LABELS[slug] for slug in PHASE_ORDER]
    rows = row_map(phase)

    fig, axes = plt.subplots(
        3,
        1,
        figsize=(7.2, 8.1),
        dpi=220,
        sharex=True,
        gridspec_kw={"height_ratios": [1.15, 1.0, 1.0], "hspace": 0.32},
    )

    for metric_key, metric_label in [("P_f1", "P F1"), ("S_f1", "S F1"), ("mean_f1", "Mean F1")]:
        values = [float(rows[slug][metric_key]) for slug in PHASE_ORDER]
        plot_line_points(axes[0], x, values, metric_label)
    axes[0].set_ylabel("F1 on unfiltered test set", fontsize=9)
    axes[0].set_ylim(0.66, 0.78)
    axes[0].legend(frameon=False, ncol=3, fontsize=8, loc="upper right")
    axes[0].set_title(
        "Phase picking: matched transfer learning, 60,837 records per condition\n"
        "Filters: SNR > 5 dB and SNR > 10 dB; evaluation: 10,000 unfiltered CREDIT-X1 test windows",
        fontsize=9.2,
        loc="left",
        pad=12,
    )
    add_panel_label(axes[0], "A")
    annotate_ci_unavailable(axes[0])

    for metric_key, metric_label in [("P_f1", "P F1"), ("S_f1", "S F1"), ("mean_f1", "Mean F1")]:
        base = float(rows["full"][metric_key])
        values = [float(rows[slug][metric_key]) - base for slug in PHASE_ORDER]
        plot_line_points(axes[1], x, values, metric_label)
    axes[1].axhline(0, color="#333333", linewidth=0.9)
    axes[1].set_ylabel("Change relative to\nfull distribution", fontsize=9)
    axes[1].set_ylim(-0.06, 0.012)
    axes[1].legend(frameon=False, ncol=3, fontsize=8, loc="lower left")
    add_panel_label(axes[1], "B")
    annotate_ci_unavailable(axes[1])

    for metric_key, metric_label, marker in [
        ("P_precision", "P precision", "o"),
        ("S_precision", "S precision", "o"),
        ("P_recall", "P recall", "^"),
        ("S_recall", "S recall", "^"),
    ]:
        base = float(rows["full"][metric_key])
        values = [float(rows[slug][metric_key]) - base for slug in PHASE_ORDER]
        plot_line_points(axes[2], x, values, metric_label, marker=marker)
    axes[2].axhline(0, color="#333333", linewidth=0.9)
    axes[2].set_ylabel("Change relative to\nfull distribution", fontsize=9)
    axes[2].set_ylim(-0.14, 0.03)
    axes[2].legend(frameon=False, ncol=2, fontsize=8, loc="lower left")
    add_panel_label(axes[2], "C")
    annotate_ci_unavailable(axes[2])

    for ax in axes:
        style_axes(ax)
        ax.set_xlim(-0.15, len(x) - 0.85)
    axes[-1].set_xticks(x)
    axes[-1].set_xticklabels(labels, fontsize=8)

    fig.subplots_adjust(left=0.16, right=0.985, top=0.91, bottom=0.08, hspace=0.34)
    fig.savefig(FIG_DIR / "fig2_phase_picking_optimized.pdf", bbox_inches="tight", facecolor="white")
    fig.savefig(FIG_DIR / "fig2_phase_picking_optimized.png", bbox_inches="tight", facecolor="white")
    plt.close(fig)
    write_csv(FIG_DIR / "fig2_phase_picking_optimized_data.csv", data)


def figure3(disp: dict, data: list[dict]) -> None:
    x = np.arange(len(DISP_ORDER))
    labels = [CONDITION_LABELS[slug] for slug in DISP_ORDER]
    rows = row_map(disp)

    fig, axes = plt.subplots(
        2,
        1,
        figsize=(7.2, 5.7),
        dpi=220,
        sharex=True,
        gridspec_kw={"height_ratios": [1.1, 1.0], "hspace": 0.30},
    )

    for metric_key, metric_label in [("val_mae", "MAE"), ("val_rmse", "RMSE")]:
        values = [float(rows[slug][metric_key]) for slug in DISP_ORDER]
        plot_line_points(axes[0], x, values, metric_label)
    axes[0].set_ylabel("Phase-velocity error (km/s)", fontsize=9)
    axes[0].set_ylim(0.043, 0.073)
    axes[0].legend(frameon=False, ncol=2, fontsize=8, loc="upper left", bbox_to_anchor=(0.065, 1.0))
    axes[0].set_title(
        "Dispersion estimation: matched from-scratch training, 11,033 samples per condition\n"
        "Filters: SNR > 3.04 dB and SNR > 6.77 dB; evaluation: 8,292 unfiltered NCF test samples",
        fontsize=9.2,
        loc="left",
        pad=12,
    )
    axes[0].text(
        0.98,
        0.84,
        "lower is better",
        transform=axes[0].transAxes,
        ha="right",
        fontsize=8,
        color="#555555",
    )
    add_panel_label(axes[0], "A")
    annotate_ci_unavailable(axes[0])

    for metric_key, metric_label in [("val_mae", "MAE"), ("val_rmse", "RMSE")]:
        base = float(rows["full"][metric_key])
        values = [float(rows[slug][metric_key]) - base for slug in DISP_ORDER]
        plot_line_points(axes[1], x, values, metric_label)
    axes[1].axhline(0, color="#333333", linewidth=0.9)
    axes[1].set_ylabel("Error increase relative to\nfull distribution (km/s)", fontsize=9)
    axes[1].set_ylim(-0.0006, 0.0060)
    axes[1].legend(frameon=False, ncol=2, fontsize=8, loc="upper left", bbox_to_anchor=(0.065, 1.0))
    add_panel_label(axes[1], "B")
    annotate_ci_unavailable(axes[1])

    for ax in axes:
        style_axes(ax)
        ax.set_xlim(-0.15, len(x) - 0.85)
    axes[-1].set_xticks(x)
    axes[-1].set_xticklabels(labels, fontsize=8)

    fig.subplots_adjust(left=0.18, right=0.985, top=0.88, bottom=0.17, hspace=0.34)
    fig.savefig(FIG_DIR / "fig3_dispersion_optimized.pdf", bbox_inches="tight", facecolor="white")
    fig.savefig(FIG_DIR / "fig3_dispersion_optimized.png", bbox_inches="tight", facecolor="white")
    plt.close(fig)
    write_csv(FIG_DIR / "fig3_dispersion_optimized_data.csv", data)


def check_manuscript_numbers(phase: dict, disp: dict) -> list[str]:
    tex_path = ROOT / "grl_overleaf" / "main.tex"
    text = tex_path.read_text(encoding="utf-8") if tex_path.exists() else ""
    expected = {
        "phase mean F1 full": f"{row_map(phase)['full']['mean_f1']:.3f}",
        "phase mean F1 moderate": f"{row_map(phase)['snr5']['mean_f1']:.3f}",
        "phase mean F1 strict": f"{row_map(phase)['snr10']['mean_f1']:.3f}",
        "phase P precision full": f"{row_map(phase)['full']['P_precision']:.3f}",
        "phase P precision strict": f"{row_map(phase)['snr10']['P_precision']:.3f}",
        "phase S precision full": f"{row_map(phase)['full']['S_precision']:.3f}",
        "phase S precision strict": f"{row_map(phase)['snr10']['S_precision']:.3f}",
        "dispersion MAE full": f"{row_map(disp)['full']['val_mae']:.4f}",
        "dispersion MAE moderate": f"{row_map(disp)['snr_q1']['val_mae']:.4f}",
        "dispersion MAE strict": f"{row_map(disp)['snr_q2']['val_mae']:.4f}",
        "dispersion RMSE strict": f"{row_map(disp)['snr_q2']['val_rmse']:.4f}",
    }
    warnings = []
    for label, value in expected.items():
        if value not in text:
            warnings.append(f"- WARNING: manuscript does not contain rounded {label} value `{value}`.")
    return warnings


def write_report(phase: dict, disp: dict, warnings: list[str]) -> None:
    phase_candidates = per_sample_candidates(PHASE_DIR)
    disp_candidates = per_sample_candidates(DISP_DIR)
    lines = [
        "# Figure Metric Discrepancy Report",
        "",
        f"Generated: {datetime.now(timezone.utc).isoformat()}",
        "",
        "## Source Files",
        "",
        f"- Phase summary: `{PHASE_DIR / 'summary.json'}`",
        f"- Dispersion summary: `{DISP_DIR / 'summary.json'}`",
        "",
        "## Per-Sample Output Scan",
        "",
        f"- Phase per-window candidates in target output directory: {phase_candidates or 'none found'}",
        f"- Dispersion per-sample candidates in target output directory: {disp_candidates or 'none found'}",
        "",
        "## CI Method Status",
        "",
        "- Requested method: paired bootstrap over shared test windows/samples.",
        "- Implemented status: unavailable for this figure build because the target experiment outputs are aggregate-only.",
        "- The optimized figures plot verified point estimates and paired deltas; error bars are omitted and marked as unavailable.",
        "",
        "## Manuscript Number Check",
        "",
    ]
    if warnings:
        lines.extend(warnings)
    else:
        lines.append("- No rounded point-estimate discrepancies detected between summary JSON values and manuscript text.")
    lines.extend(
        [
            "",
            "## Plotted Point Estimates",
            "",
            f"- Phase mean F1: full={row_map(phase)['full']['mean_f1']:.6f}, "
            f"moderate={row_map(phase)['snr5']['mean_f1']:.6f}, "
            f"strict={row_map(phase)['snr10']['mean_f1']:.6f}.",
            f"- Dispersion MAE: full={row_map(disp)['full']['val_mae']:.6f}, "
            f"moderate={row_map(disp)['snr_q1']['val_mae']:.6f}, "
            f"strict={row_map(disp)['snr_q2']['val_mae']:.6f}.",
            "",
        ]
    )
    REPORT.parent.mkdir(parents=True, exist_ok=True)
    REPORT.write_text("\n".join(lines), encoding="utf-8")


def main() -> None:
    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 8.5,
            "axes.titlesize": 10,
            "axes.labelsize": 9,
            "legend.fontsize": 8,
            "pdf.fonttype": 42,
            "ps.fonttype": 42,
            "savefig.transparent": False,
        }
    )
    FIG_DIR.mkdir(parents=True, exist_ok=True)
    phase = read_json(PHASE_DIR / "summary.json")
    disp = read_json(DISP_DIR / "summary.json")
    phase_data = make_phase_data(phase)
    disp_data = make_disp_data(disp)
    figure2(phase, phase_data)
    figure3(disp, disp_data)
    warnings = check_manuscript_numbers(phase, disp)
    write_report(phase, disp, warnings)
    print(f"Wrote optimized figures to {FIG_DIR}")
    print(f"Wrote discrepancy report to {REPORT}")


if __name__ == "__main__":
    main()