"""Journal-style figures (vector PDF + PNG) drawn only from saved result files. Times/STIX typography, monochrome
marks distinguished by shape and line style, one restrained accent colour, panel labels, inward ticks."""
from __future__ import annotations

import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402

INK, INK2, INK3, INK4 = "#000000", "#4d4d4d", "#8c8c8c", "#c8c8c8"
ACCENT = "#1f4e79"
plt.rcParams.update({
    "font.family": "serif", "font.serif": ["Times New Roman", "Times", "STIX Two Text", "DejaVu Serif"], "mathtext.fontset": "stix",
    "font.size": 9, "axes.titlesize": 9, "axes.labelsize": 9, "legend.fontsize": 7.5, "xtick.labelsize": 8, "ytick.labelsize": 8,
    "axes.spines.top": False, "axes.spines.right": False, "axes.edgecolor": INK, "axes.linewidth": 0.6, "xtick.direction": "in",
    "ytick.direction": "in", "xtick.major.size": 3, "ytick.major.size": 3, "axes.grid": False, "legend.frameon": False,
    "lines.linewidth": 1.0, "lines.markersize": 4, "lines.markeredgewidth": 0.8, "figure.dpi": 150, "savefig.bbox": "tight",
    "savefig.pad_inches": 0.02, "pdf.fonttype": 42})
W1, W2 = 3.25, 6.5
STYLE = {"naive_rw": dict(marker="o", ls=":", color=INK3), "seasonal_vol": dict(marker="v", ls=":", color=INK3), "ewma": dict(marker="^", ls=":", color=INK3),
         "har_rv": dict(marker="s", ls="--", color=INK2), "garch_t": dict(marker="D", ls="--", color=INK2), "ridge": dict(marker="x", ls="-.", color=INK3),
         "elastic_net": dict(marker="+", ls="-.", color=INK3), "lightgbm": dict(marker="p", ls="-.", color=INK2), "mlp": dict(marker="<", ls="-", color=INK3),
         "tcn": dict(marker=">", ls="-", color=INK2), "mrpt_small_1m_only": dict(marker="s", ls="-", color=INK), "mrpt_small": dict(marker="o", ls="-", color=ACCENT),
         "mrpt_base": dict(marker="D", ls="-", color=ACCENT), "mrpt_base_1m_only": dict(marker="D", ls="--", color=INK), "mrpt_spec": dict(marker="h", ls="-", color=ACCENT)}
LABEL = {"naive_rw": "Random walk", "seasonal_vol": "Seasonal vol.", "ewma": "EWMA-$t$", "har_rv": "HAR-RV", "garch_t": "GARCH(1,1)-$t$", "ridge": "Ridge",
         "elastic_net": "Elastic net", "lightgbm": "LightGBM", "mlp": "MLP (1m)", "tcn": "TCN (1m)", "mrpt_small_1m_only": "MRPT-S 1m only",
         "mrpt_small": "MRPT-S multi-res.", "mrpt_base": "MRPT-B multi-res.", "mrpt_base_1m_only": "MRPT-B 1m only", "mrpt_spec": "MRPT-M multi-res.",
         "mrpt_small_b4k": "MRPT-S ref. (4k)", "mrpt_small_1m_only_b4k": "1m only (4k)", "mrpt_small_1m_5m": "1m+5m", "mrpt_small_cross": "cross-attention fusion",
         "mrpt_small_gaussian": "Gaussian head", "mrpt_small_noregime": "no regime", "mrpt_small_learnedregime": "learned regime", "mrpt_small_returns_only": "returns only",
         "mrpt_small_time_only": "time only", "mrpt_small_ohlcv": "no relative volume", "mrpt_small_shuffled": "shuffled minutes", "mrpt_small_ticker": "ticker embedding",
         "mrpt_small_seed2": "ref. seed 2", "mrpt_small_seed3": "ref. seed 3", "mrpt_small_copula": "copula, joint", "mrpt_small_copula_indep": "copula, independent"}


def _panel(ax, letter, title):
    ax.set_title(f"Panel {letter}. {title}", loc="left", fontsize=9, fontweight="bold", pad=6)


def _save(fig, out: Path, name: str):
    out.mkdir(parents=True, exist_ok=True)
    fig.savefig(out / f"{name}.pdf"); fig.savefig(out / f"{name}.png", dpi=220); plt.close(fig)


def _lab(m):
    base = m.replace("_heldout", "")
    return LABEL.get(base, base.replace("_", " ")) + (" (held-out)" if m.endswith("_heldout") else "")


def fig_scores_by_horizon(lb: pd.DataFrame, out: Path, name: str, metric: str = "crps_scaled", ref: str = "naive_rw", models=None):
    """Panel A: metric by horizon; Panel B: relative change vs the reference with bootstrap CIs."""
    lb = lb.copy(); lb["base"] = lb["model"].str.replace("_heldout", "", regex=False)
    models = models or [m for m in STYLE if m in set(lb["base"])]
    H = sorted(lb["H"].unique())
    fig, axes = plt.subplots(1, 2, figsize=(W2, 2.6))
    for m in models:
        d = lb[lb["base"] == m].sort_values("H")
        if d.empty:
            continue
        st = STYLE.get(m, dict(marker=".", ls="-", color=INK3))
        axes[0].plot(d["H"], d[metric], label=_lab(m), **st)
        r = lb[lb["base"] == ref].set_index("H")[metric]
        rel = 100 * (d.set_index("H")[metric] / r.loc[d["H"]].values - 1)
        axes[1].plot(d["H"], rel.values, label=_lab(m), **st)
    _panel(axes[0], "A", f"{metric.replace('_', ' ')} by horizon"); axes[0].set_xlabel("Horizon $H$ (minutes)"); axes[0].set_ylabel(metric.replace("_", " "))
    _panel(axes[1], "B", f"relative to {_lab(ref)} (%)"); axes[1].axhline(0, color=INK4, lw=0.6); axes[1].set_xlabel("Horizon $H$ (minutes)"); axes[1].set_ylabel("%")
    for ax in axes:
        ax.set_xticks(H)
    h_, l_ = axes[1].get_legend_handles_labels()
    fig.legend(h_, l_, loc="lower center", ncol=4, bbox_to_anchor=(0.5, -0.36), frameon=False)
    fig.subplots_adjust(wspace=0.3)
    _save(fig, out, name)


def fig_calibration(lb: pd.DataFrame, out: Path, name: str, models=None):
    """Coverage of 50/80/95 % central intervals at each horizon (nominal = dotted)."""
    lb = lb.copy(); lb["base"] = lb["model"].str.replace("_heldout", "", regex=False)
    models = models or [m for m in STYLE if m in set(lb["base"])]
    H = sorted(lb["H"].unique())
    fig, axes = plt.subplots(1, 3, figsize=(W2, 2.4), sharey=False)
    for ax, lv in zip(axes, (50, 80, 95)):
        for m in models:
            d = lb[lb["base"] == m].sort_values("H")
            if d.empty:
                continue
            ax.plot(d["H"], 100 * d[f"cov_{lv}"], label=_lab(m), **STYLE.get(m, dict(marker=".", ls="-", color=INK3)))
        ax.axhline(lv, color=INK4, lw=0.8, ls=":"); ax.set_xticks(H); ax.set_xlabel("$H$ (min)")
        _panel(ax, "ABC"[(50, 80, 95).index(lv)], f"{lv}% coverage")
    h_, l_ = axes[0].get_legend_handles_labels()
    fig.legend(h_, l_, loc="lower center", ncol=4, bbox_to_anchor=(0.5, -0.34), frameon=False)
    fig.subplots_adjust(wspace=0.35)
    _save(fig, out, name)


def fig_pit(per: dict[str, np.ndarray], out: Path, name: str, H: int = 60, bins: int = 20):
    """PIT histograms of selected models at one horizon."""
    n = len(per)
    fig, axes = plt.subplots(1, n, figsize=(W2, 1.9), sharey=True)
    axes = np.atleast_1d(axes)
    for ax, (m, u), letter in zip(axes, per.items(), "ABCDEFGH"):
        ax.hist(u, bins=bins, range=(0, 1), density=True, color=INK4, edgecolor=INK, lw=0.4)
        ax.axhline(1, color=ACCENT, lw=0.8, ls="--"); ax.set_xlim(0, 1); ax.set_xlabel("PIT")
        _panel(ax, letter, f"{_lab(m)}, $H={H}$")
    axes[0].set_ylabel("density")
    _save(fig, out, name)


def fig_ablation_forest(cmp: pd.DataFrame, out: Path, name: str, ref_label: str, H: int = 60, order=None, labels=None):
    """Relative CRPS change versus a reference with day-cluster bootstrap CIs (forest plot) at one horizon."""
    d = cmp[cmp["H"] == H].copy()
    if order:
        d = d.set_index("model").loc[[o for o in order if o in d["model"].values]].reset_index()
    y = np.arange(len(d))[::-1]
    ref_mean = (d["diff"] / d["rel_diff"]).replace(0, np.nan)          # mean reference loss (diff = rel_diff * ref_mean)
    rel, lo, hi = 100 * d["rel_diff"], 100 * d["lo"] / ref_mean, 100 * d["hi"] / ref_mean
    fig, ax = plt.subplots(figsize=(W2, 0.28 * len(d) + 0.9))
    ax.errorbar(rel, y, xerr=[(rel - lo).abs().fillna(0), (hi - rel).abs().fillna(0)], fmt="o", color=INK, ecolor=INK2, capsize=2.5,
                elinewidth=0.7, capthick=0.7, ms=4)
    ax.axvline(0, color=INK4, lw=0.8)
    ax.set_yticks(y); ax.set_yticklabels([(labels or {}).get(m, _lab(m)) for m in d["model"]])
    ax.set_xlabel(f"CRPS change vs {ref_label} at $H={H}$ (%), 95% day-cluster bootstrap CI")
    _save(fig, out, name)


def fig_training_curves(logs: dict[str, Path], out: Path, name: str):
    fig, axes = plt.subplots(1, 2, figsize=(W2, 2.4))
    for (m, p), st in zip(logs.items(), [dict(ls="-", color=ACCENT), dict(ls="--", color=INK), dict(ls=":", color=INK2), dict(ls="-.", color=INK3)] * 3):
        rows = [json.loads(l) for l in Path(p).read_text().splitlines() if l.strip()]
        tr = pd.DataFrame([r for r in rows if r.get("event") == "train"]); va = pd.DataFrame([r for r in rows if r.get("event") == "val"])
        if len(tr):
            axes[0].plot(tr["step"], tr["loss"].rolling(10, min_periods=1).mean(), label=_lab(m), lw=0.9, **st)
        if len(va):
            axes[1].plot(va["step"], va["val_nll"], marker="o", ms=3, label=_lab(m), **st)
    _panel(axes[0], "A", "training NLL (moving average)"); _panel(axes[1], "B", "validation NLL")
    for ax in axes:
        ax.set_xlabel("optimiser step")
    axes[1].legend(loc="best")
    _save(fig, out, name)


def fig_subgroups(df: pd.DataFrame, out: Path, name: str, col: str, metric_label: str):
    """Relative CRPS change of the main model vs reference by subgroup (bars with CIs)."""
    y = np.arange(len(df))[::-1]
    fig, ax = plt.subplots(figsize=(W1 * 1.6, 0.26 * len(df) + 0.9))
    ax.barh(y, 100 * df["rel"], color=INK4, edgecolor=INK, lw=0.5, height=0.6)
    ax.errorbar(100 * df["rel"], y, xerr=[100 * (df["rel"] - df["lo"]), 100 * (df["hi"] - df["rel"])], fmt="none", ecolor=INK, elinewidth=0.7, capsize=2)
    ax.axvline(0, color=INK2, lw=0.8); ax.set_yticks(y); ax.set_yticklabels([str(v) for v in df[col]]); ax.set_xlabel(metric_label)
    _save(fig, out, name)