#!/usr/bin/env python3
"""Create the matched SNR filtering design figure for the GRL manuscript."""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

import matplotlib.pyplot as plt


ROOT = Path(__file__).resolve().parents[1]
PHASE_SUMMARY = ROOT / "outputs" / "snr_transfer_seed20260609" / "summary.json"
DISP_SUMMARY = ROOT / "outputs" / "disp_snr_transfer_seed20260609" / "summary.json"
DISP_SNR_CACHE = ROOT / "outputs" / "disp_snr_transfer_seed20260609" / "ncf_snr_cache.json"
DEFAULT_OUTPUT = ROOT / "grl_overleaf" / "figures" / "snr_matched_design.png"


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


def dispersion_candidate_counts(summary: dict) -> list[int]:
    cache = read_json(DISP_SNR_CACHE)
    q1 = summary["snr_thresholds"]["q1"]
    q2 = summary["snr_thresholds"]["q2"]
    train_snr = [
        item["snr_db"]
        for item in cache.values()
        if item.get("split") == "train" and math.isfinite(float(item.get("snr_db", float("nan"))))
    ]
    return [
        len(train_snr),
        sum(float(value) > q1 for value in train_snr),
        sum(float(value) > q2 for value in train_snr),
    ]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    args = parser.parse_args()

    phase = read_json(PHASE_SUMMARY)
    disp = read_json(DISP_SUMMARY)

    phase_candidate = [
        phase["candidate_train_records"]["full"],
        phase["candidate_train_records"]["snr5"],
        phase["candidate_train_records"]["snr10"],
    ]
    phase_matched = [row["train_records"] for row in phase["rows"]]
    disp_matched = [row["train_records"] for row in disp["rows"]]
    disp_candidate = dispersion_candidate_counts(disp)

    labels = ["Full", "SNR > 5", "SNR > 10"]
    disp_labels = ["Full", "SNR > 3.04", "SNR > 6.77"]
    colors = ["#2f6f73", "#c47f2b", "#8f4d6b"]

    fig, axes = plt.subplots(1, 2, figsize=(10.8, 4.6), dpi=220)
    panels = [
        (axes[0], "A", "Phase picking", labels, phase_candidate, phase_matched, "records"),
        (axes[1], "B", "Dispersion", disp_labels, disp_candidate, disp_matched, "samples"),
    ]
    for ax, panel_label, title, xlabels, candidate, matched, unit in panels:
        x = range(len(xlabels))
        ax.bar(x, candidate, width=0.62, color="#d7d7d7", edgecolor="#555555", label="candidate pool")
        ax.bar(x, matched, width=0.38, color=colors, edgecolor="#222222", label="matched training")
        for idx, value in enumerate(candidate):
            ax.text(idx, value * 1.02, f"{value:,}", ha="center", va="bottom", fontsize=7)
        for idx, value in enumerate(matched):
            ax.text(idx, value * 0.50, f"{value:,}", ha="center", va="center", fontsize=8, color="white")
        ax.set_title(title, fontsize=11)
        ax.set_xticks(list(x))
        ax.set_xticklabels(xlabels, fontsize=8)
        ax.set_ylabel(f"training {unit}")
        ax.spines[["top", "right"]].set_visible(False)
        ax.grid(axis="y", color="#eeeeee", linewidth=0.8)
        ax.set_axisbelow(True)
        ax.text(
            -0.12,
            1.08,
            panel_label,
            transform=ax.transAxes,
            fontsize=13,
            fontweight="bold",
            va="top",
        )
    axes[0].legend(frameon=False, loc="upper right", fontsize=8)
    fig.suptitle("Matched training budgets isolate the effect of SNR filtering", fontsize=13, y=0.98)
    fig.tight_layout(rect=(0, 0, 1, 0.94))

    args.output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(args.output, bbox_inches="tight")
    print(f"Wrote {args.output}")


if __name__ == "__main__":
    main()