"""Journal figures (vector PDF, matplotlib). One axis per panel, thin marks, fixed categorical order, direct labels.
Palette validated with the dataviz validator (light surface): blue, orange, aqua, yellow (+ gray for references)."""
from __future__ import annotations

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

SERIES = {"daily": "#2a78d6", "30m": "#eb6834", "5m": "#1baf7a", "1m": "#eda100", "fund": "#6250d6", "ref": "#8a8a86", "neg": "#e34948"}
MODEL_COLOR = {"B1": SERIES["fund"], "B2": SERIES["daily"], "B3": SERIES["30m"], "B4": SERIES["5m"], "B5": SERIES["1m"], "B5all": "#52514e",
               "LB1": SERIES["fund"], "LB2": SERIES["daily"], "LB3": SERIES["30m"], "LB4": SERIES["5m"], "LB5": SERIES["1m"], "LB5all": "#52514e"}
MODEL_LABEL = {"B1": "B1 fundamentals", "B2": "B2 + daily", "B3": "B3 + 30-min", "B4": "B4 + 5-min", "B5": "B5 + 1-min", "B5all": "all frequencies"}

plt.rcParams.update({"font.family": "serif", "font.size": 9, "axes.spines.top": False, "axes.spines.right": False,
                     "axes.edgecolor": "#8a8a86", "axes.linewidth": 0.6, "xtick.color": "#52514e", "ytick.color": "#52514e",
                     "axes.labelcolor": "#0b0b0b", "grid.color": "#e6e5e1", "grid.linewidth": 0.5, "legend.frameon": False,
                     "figure.dpi": 150, "savefig.bbox": "tight"})


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=200); plt.close(fig)


def fig_timeline(out: Path, h_list=(40, 20, 10, 5, 1)):
    """Event-time timeline and the no-leakage rule (schematic drawn from the protocol, labelled as such)."""
    fig, ax = plt.subplots(figsize=(6.4, 2.2))
    ax.set_xlim(-75, 15); ax.set_ylim(0, 3); ax.axis("off")
    ax.plot([-75, 12], [1, 1], color="#52514e", lw=1)
    for x, lab, c in ((-63, "$A_{q-1}$\nprevious release", SERIES["ref"]), (0, "$A_q$\nrelease of quarter $q$", SERIES["neg"]),
                      (1, "$F_q$ filing", SERIES["ref"]), (2, "$V_q=F_q+1$\nfacts available", SERIES["ref"])):
        ax.plot([x], [1], marker="|", ms=14, color=c, mew=1.5)
        ax.text(x, 0.15 if x not in (1,) else 1.55, lab, ha="center", va="center", fontsize=7, color="#0b0b0b")
    for i, h in enumerate(h_list):
        ax.plot([-h], [1], marker="o", ms=4, color=SERIES["1m"], mec="white", mew=0.8)
        ax.text(-h, 1.35 + 0.28 * (i % 2), f"$c=A_q-{h}$", ha="center", fontsize=6.5, color="#52514e")
    ax.add_patch(plt.Rectangle((-40, 0.85), 20, 0.3, color=SERIES["1m"], alpha=0.25, lw=0))
    ax.text(-30, 2.6, "20-session feature window ends at $c$; bars after $c$ and facts with $V_j>c$ never enter the information set",
            ha="center", fontsize=7, color="#52514e")
    ax.set_title("Event-time design (schematic of the protocol, not data)", fontsize=8, loc="left")
    _save(fig, out, "fig01_timeline")


def fig_coverage(out: Path, years: pd.DataFrame, sizes: pd.DataFrame | None = None):
    """years: columns year, firms, firm_quarters, adverse_rate."""
    fig, axes = plt.subplots(1, 2, figsize=(6.4, 2.4))
    ax = axes[0]
    ax.bar(years["year"], years["firms"], color=SERIES["daily"], width=0.8)
    ax.set_ylabel("Firms in sample"); ax.set_xlabel("Year of cutoff"); ax.grid(axis="y")
    for x, v in zip(years["year"].iloc[::3], years["firms"].iloc[::3]):
        ax.text(x, v, f"{int(v):,}", ha="center", va="bottom", fontsize=6, color="#52514e")
    ax = axes[1]
    ax.plot(years["year"], years["adverse_rate"], color=SERIES["neg"], lw=2, marker="o", ms=4, mec="white")
    ax.set_ylabel("Adverse-surprise rate"); ax.set_xlabel("Year of cutoff"); ax.grid(axis="y"); ax.set_ylim(0, None)
    _save(fig, out, "fig02_coverage")


def fig_prevalence(out: Path, prev: pd.DataFrame):
    """prev: columns group (str), label, rate, lo, hi  — conditional adverse prevalence by ex-ante characteristics."""
    fig, ax = plt.subplots(figsize=(6.0, 2.6))
    y = np.arange(len(prev))
    ax.hlines(y, prev["lo"], prev["hi"], color="#c3c2b7", lw=2)
    ax.plot(prev["rate"], y, "o", color=SERIES["neg"], ms=5, mec="white")
    ax.set_yticks(y); ax.set_yticklabels(prev["label"]); ax.invert_yaxis()
    ax.set_xlabel("Share of firm-quarters with an adverse operating-income surprise (test period)"); ax.grid(axis="x")
    _save(fig, out, "fig03_prevalence")


def fig_frequency_gains(out: Path, contrasts: pd.DataFrame):
    """Paired Δ log loss vs B2 for B3/B4/B5 (logistic and LightGBM), with 95 % block-bootstrap intervals."""
    fam = [("logistic", ""), ("lightgbm", "L")]
    fig, axes = plt.subplots(1, 2, figsize=(6.4, 2.6), sharey=True)
    for ax, (name, pre) in zip(axes, fam):
        sub = contrasts[(contrasts.family == name) & (contrasts.b == pre + "B2") & (contrasts.a.isin([pre + "B3", pre + "B4", pre + "B5", pre + "B5all"]))]
        order = [pre + "B3", pre + "B4", pre + "B5", pre + "B5all"]
        sub = sub.set_index("a").reindex([o for o in order if o in sub.a.values]).reset_index()
        x = np.arange(len(sub))
        for i, r in sub.iterrows():
            c = MODEL_COLOR.get(r["a"].replace("L", "", 1) if r["a"].startswith("L") else r["a"], "#52514e")
            ax.vlines(i, r["ci_lo"], r["ci_hi"], color=c, lw=2)
            ax.plot(i, r["d_logloss"], "o", color=c, ms=6, mec="white")
            ax.text(i, r["ci_hi"], f"{r['d_logloss']:+.4f}", ha="center", va="bottom", fontsize=6.5, color="#52514e")
        ax.axhline(0, color="#8a8a86", lw=0.8)
        ax.set_xticks(x); ax.set_xticklabels([MODEL_LABEL[a.replace("L", "", 1) if a.startswith("L") else a].replace("B5all", "all") for a in sub.a], fontsize=7)
        ax.set_title(f"{name}", fontsize=8, loc="left"); ax.grid(axis="y")
    axes[0].set_ylabel("Δ log loss vs. B2 (daily + accounting)\nnegative = better")
    _save(fig, out, "fig04_frequency_gains")


def fig_reliability(out: Path, rel: pd.DataFrame):
    fig, ax = plt.subplots(figsize=(3.6, 3.2))
    ax.plot([0, rel[["mean_pred", "obs_rate"]].max().max()], [0, rel[["mean_pred", "obs_rate"]].max().max()], color="#8a8a86", lw=0.8, ls="--")
    for m in ("LB2", "LB5", "B2", "B5"):
        s = rel[rel.model == m]
        if len(s):
            ax.plot(s["mean_pred"], s["obs_rate"], marker="o", ms=4, mec="white", lw=1.6, color=MODEL_COLOR[m], label=m)
    ax.set_xlabel("Mean predicted probability (decile)"); ax.set_ylabel("Observed adverse rate"); ax.legend(fontsize=7); ax.grid()
    _save(fig, out, "fig05_reliability")


def fig_signal_evolution(out: Path, ladder: pd.DataFrame):
    """Mean AUC and Δ log loss by horizon (all firm-quarters, no selection)."""
    fig, axes = plt.subplots(1, 2, figsize=(6.4, 2.5))
    for fam, c in (("logistic", SERIES["daily"]), ("lightgbm", SERIES["1m"])):
        s = ladder[ladder.family == fam].sort_values("h", ascending=False)
        axes[0].plot(-s["h"], s["auc_a"], marker="o", ms=4, mec="white", lw=1.6, color=c, label=f"{fam}: B5")
        axes[0].plot(-s["h"], s["auc_b"], marker="s", ms=3.5, mec="white", lw=1.2, ls="--", color=c, label=f"{fam}: B2")
        axes[1].errorbar(-s["h"], s["d_logloss"], yerr=[s["d_logloss"] - s["ci_lo"], s["ci_hi"] - s["d_logloss"]], fmt="o-", ms=4, lw=1.4,
                         color=c, capsize=2, label=fam)
    axes[0].set_xlabel("Trading days before the release (−h)"); axes[0].set_ylabel("Test AUC"); axes[0].legend(fontsize=6); axes[0].grid()
    axes[1].axhline(0, color="#8a8a86", lw=0.8); axes[1].set_xlabel("Trading days before the release (−h)")
    axes[1].set_ylabel("Δ log loss B5 − B2"); axes[1].legend(fontsize=7); axes[1].grid()
    _save(fig, out, "fig06_signal_evolution")


def fig_heterogeneity(out: Path, sub: pd.DataFrame, groups: list[str]):
    s = sub[(sub.family == "lightgbm") & sub.subgroup.isin(groups)].set_index("subgroup").reindex(groups).reset_index()
    fig, ax = plt.subplots(figsize=(6.0, 2.8))
    y = np.arange(len(s))
    ax.hlines(y, s["ci_lo"], s["ci_hi"], color="#c3c2b7", lw=2)
    ax.plot(s["mean"], y, "o", color=SERIES["1m"], ms=5, mec="white")
    for yi, r in zip(y, s.itertuples()):
        ax.text(r.ci_hi, yi, f"  n={int(r.n):,}", va="center", fontsize=6, color="#52514e")
    ax.axvline(0, color="#8a8a86", lw=0.8)
    ax.set_yticks(y); ax.set_yticklabels([g.replace("_", " ") for g in s["subgroup"]]); ax.invert_yaxis()
    ax.set_xlabel("Δ log loss LB5 − LB2 (test period), 95 % block-bootstrap CI"); ax.grid(axis="x")
    _save(fig, out, "fig07_heterogeneity")


def fig_rolling(out: Path, wf: pd.DataFrame, wfc: pd.DataFrame):
    fig, axes = plt.subplots(1, 2, figsize=(6.4, 2.5))
    for m in ("LB2", "LB5", "B2", "B5"):
        s = wf[wf.model == m].sort_values("year")
        if len(s):
            axes[0].plot(s["year"], s["auc"], marker="o", ms=3.5, mec="white", lw=1.4, color=MODEL_COLOR[m], ls="-" if m.startswith("L") else "--", label=m)
    axes[0].set_ylabel("Out-of-fold AUC (annual refits)"); axes[0].legend(fontsize=6, ncol=2); axes[0].grid()
    for fam, c in (("lightgbm", SERIES["1m"]), ("logistic", SERIES["daily"])):
        s = wfc[(wfc.family == fam) & (wfc.year != "all")].copy(); s["year"] = s["year"].astype(int); s = s.sort_values("year")
        axes[1].errorbar(s["year"], s["d_logloss"], yerr=[s["d_logloss"] - s["ci_lo"], s["ci_hi"] - s["d_logloss"]], fmt="o-", ms=3.5, lw=1.2, capsize=2, color=c, label=fam)
    axes[1].axhline(0, color="#8a8a86", lw=0.8); axes[1].set_ylabel("Δ log loss B5 − B2 by year"); axes[1].legend(fontsize=7); axes[1].grid()
    _save(fig, out, "fig08_rolling")


def fig_economic(out: Path, port: pd.DataFrame, alphas: dict):
    fig, ax = plt.subplots(figsize=(4.2, 2.6))
    dm = port.groupby("decile")["ret"].mean() * 100
    ax.bar(dm.index, dm.values, color=SERIES["daily"], width=0.8)
    ax.axhline(0, color="#8a8a86", lw=0.8); ax.set_xlabel("Decile of ex-ante deterioration score (10 = most likely)")
    ax.set_ylabel("Mean monthly return (%)"); ax.grid(axis="y")
    ax.set_title(f"L/S 1−10 alpha: {alphas['gross']['alpha_monthly']*100:+.2f}% (t={alphas['gross']['alpha_t']:.1f}) gross, "
                 f"{alphas['net20bps']['alpha_monthly']*100:+.2f}% net of 20 bps", fontsize=7, loc="left")
    _save(fig, out, "fig09_economic")