"""Plot Fig 7 (XGBoost pre-transition attribution) from ../data/fig7.pkl."""
import pickle
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
import _style as S

HERE = Path(__file__).resolve().parent
D = pickle.load(open(HERE.parent / "data" / "fig7.pkl", "rb"))["results"]
OUT = HERE.parent / "figures" / "Fig7_xgboost_attribution.pdf"
C_WARM, C_COOL, C_BAR = "#d62728", "#1f77b4", "#2ca02c"
COLORS = [S.CLEAR, "#e377c2", "#8c564b", C_BAR, C_BAR]


def waterfall(ax, dr, title, color, letter, show_y):
    rows = dr["rows"]; labels = [r[0] for r in rows]
    r2 = [100 * r[1] if np.isfinite(r[1]) else 0 for r in rows]
    marg = [100 * r[2] if np.isfinite(r[2]) else 0 for r in rows]
    lo = [100 * r[3] if np.isfinite(r[3]) else 0 for r in rows]
    hi = [100 * r[4] if np.isfinite(r[4]) else 0 for r in rows]
    xs = np.arange(len(rows)); bottoms = [0.0] + [r2[i] - marg[i] for i in range(1, len(rows))]
    for i in range(len(rows)):
        ax.bar(xs[i], marg[i], bottom=bottoms[i], width=0.62, color=COLORS[i], edgecolor="white", linewidth=0.5)
        if r2[i] > 0:
            ax.errorbar(xs[i], r2[i], yerr=[[max(r2[i] - lo[i], 0)], [max(hi[i] - r2[i], 0)]],
                        fmt="none", ecolor="0.3", elinewidth=0.7, capsize=2.5, capthick=0.7, zorder=5)
            top = max(hi[i], r2[i]); close = i > 0 and r2[i - 1] > 0 and abs(r2[i] - r2[i - 1]) < 2.0
            ax.text(i, top + (1.8 if close else 0.8), f"{r2[i]:.1f}%", ha="center", va="bottom", fontsize=6)
    ax.set_xticks(xs); ax.set_xticklabels(labels, fontsize=6, rotation=30, ha="right")
    ax.set_xlabel("Cumulative predictors", fontsize=8)
    if show_y: ax.set_ylabel(r"LOSO $R^2$ (%)", fontsize=8)
    else: plt.setp(ax.get_yticklabels(), visible=False)
    ax.set_ylim(bottom=0); ax.set_title(title, fontsize=8, fontweight="bold", color=color)
    ax.text(-0.12 if show_y else -0.06, 1.10, letter, transform=ax.transAxes, fontsize=11, fontweight="bold", va="top")


def main():
    figw = 15.0 / 2.54
    fig = plt.figure(figsize=(figw, figw * 0.38)); gs = fig.add_gridspec(1, 2, wspace=0.15)
    aa = fig.add_subplot(gs[0, 0]); ab = fig.add_subplot(gs[0, 1], sharey=aa)
    cc = D["Clear → Cloudy"]; wc = D["Cloudy → Clear"]
    waterfall(aa, cc, f"Clear → Cloudy ($n$ = {cc['n_cross']})", C_WARM, "A", True)
    waterfall(ab, wc, f"Cloudy → Clear ($n$ = {wc['n_cross']})", C_COOL, "B", False)
    aa.set_ylim(0, 42)
    w, h = S.finalize(fig, OUT); print(f"wrote {OUT.name} ({w:.1f}x{h:.1f}cm)")


if __name__ == "__main__":
    main()