"""Generate paper figures from saved results (no proprietary data). Usage: make_figures.py --exps pilot,main,ablations"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from chronos_x.eval import figures as F  # noqa: E402

OUT = ROOT / "paper" / "figures"


def main() -> None:
    ap = argparse.ArgumentParser(); ap.add_argument("--exps", default="pilot"); ap.add_argument("--split", default="test")
    a = ap.parse_args()
    made = []
    for exp in a.exps.split(","):
        d = ROOT / "results" / "metrics" / exp
        for suffix in ("", "_heldout"):
            f = d / f"leaderboard_{a.split}{suffix}.csv"
            if not f.exists():
                continue
            lb = pd.read_csv(f)
            F.fig_scores_by_horizon(lb, OUT, f"fig_{exp}_crps{suffix}", "crps_scaled"); made.append(f"fig_{exp}_crps{suffix}")
            F.fig_scores_by_horizon(lb, OUT, f"fig_{exp}_wis{suffix}", "wis"); made.append(f"fig_{exp}_wis{suffix}")
            F.fig_calibration(lb, OUT, f"fig_{exp}_calibration{suffix}"); made.append(f"fig_{exp}_calibration{suffix}")
            # PIT histograms at H=60 for up to four models
            per = {}
            for m in ("naive_rw", "garch_t", "lightgbm", "mrpt_small", "mrpt_base"):
                p = d / f"per_origin_{m}_{a.split}{suffix}.npz"
                if p.exists():
                    per[m] = np.load(p)["pit_60"]
            if per:
                F.fig_pit(dict(list(per.items())[:4]), OUT, f"fig_{exp}_pit{suffix}"); made.append(f"fig_{exp}_pit{suffix}")
            fc = d / f"compare_crps_vs_naive_rw_{a.split}{suffix}.csv"
            if fc.exists():
                F.fig_ablation_forest(pd.read_csv(fc), OUT, f"fig_{exp}_forest{suffix}", "random walk"); made.append(f"fig_{exp}_forest{suffix}")
        logs = {m: d / m / "train_log.jsonl" for m in ("mrpt_small", "mrpt_small_1m_only", "mrpt_base", "tcn") if (d / m / "train_log.jsonl").exists()}
        if logs:
            F.fig_training_curves(logs, OUT, f"fig_{exp}_training"); made.append(f"fig_{exp}_training")
    print("figures:", made)


if __name__ == "__main__":
    main()