"""Plot Fig 5 (seasonal decomposition) from ../data/fig5.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" / "fig5.pkl", "rb"))
OUT = HERE.parent / "figures" / "Fig5_seasonal_decomposition.pdf"
SEASONS = ("DJF", "MAM", "JJA", "SON"); sea = D["seasonal"]; C_FIT = "#666666"


def integer_shares(s_lw, s_sw, r2):
    total = max(int(round(100 * r2)), 0); sw = min(int(s_sw + 0.5), total)
    return total - sw, sw


def main():
    fig, ax = plt.subplots(2, 2, figsize=(13.6 / 2.54, 13.6 / 2.54 * 0.75))
    j = sea["JJA"]
    # A
    a = ax[0, 0]
    a.plot(j["x"], j["D_obs"], "-", lw=1.0, color=S.OBS, label=r"$D(x)$ Obs")
    a.plot(D["era5_jja_x"], D["era5_jja_D"], "-", lw=1.0, color=S.ERA5, label=r"$D(x)$ ERA5")
    a.set_xlabel(r"$F_{\mathrm{LW,net}}$ (W m$^{-2}$)"); a.set_ylabel(r"$D(x)$ ((W m$^{-2}$)$^2$ h$^{-1}$)")
    a.set_xlim(-100, 20); top = max(np.nanmax(j["D_obs"]), np.nanmax(D["era5_jja_D"])) * 1.1
    a.set_ylim(0, max(200, int(np.ceil(top / 50) * 50))); a.legend(loc="upper left", fontsize=6)
    a.text(-0.06, 1.05, "A", transform=a.transAxes, fontsize=9, fontweight="bold")
    # B
    a = ax[0, 1]
    a.plot(D["panelB_x"], D["panelB_sig"], "o-", color=S.ERA5, lw=1.0, markersize=4, markerfacecolor="white", markeredgewidth=1.0)
    a.axhline(np.mean(D["panelB_sig"]), color="gray", ls="--", lw=0.8, alpha=0.7)
    a.set_xlabel(r"$F_{\mathrm{LW,net}}$ (W m$^{-2}$)"); a.set_ylabel(r"$\sigma_{\mathcal{T}}$")
    a.set_xlim(-120, 10); a.set_ylim(0, 0.25); a.set_title("JJA cloud variability", fontsize=8)
    a.text(-0.06, 1.05, "B", transform=a.transAxes, fontsize=9, fontweight="bold")
    # C
    a = ax[1, 0]
    dlw = np.maximum(j["k_lw"] * j["g2_lw"], 0); dsw = np.maximum(j["k_sw"] * j["g2_sw"], 0)
    ok = np.isfinite(dlw) & np.isfinite(dsw)
    f1 = a.fill_between(j["x"][ok], 0, dlw[ok], alpha=0.5, color=S.C_LWP, label=r"$D_{\mathrm{LWP}}$")
    f2 = a.fill_between(j["x"][ok], dlw[ok], (dlw + dsw)[ok], alpha=0.5, color=S.C_SW, label=r"$D_{\mathrm{SW}}$")
    l1, = a.plot(j["x"], j["D_obs"], "k-", lw=1.0, label=r"$D_{\mathrm{obs}}$")
    l2, = a.plot(j["x"], j["fitted"], "--", lw=1.0, color=C_FIT, label=r"$D_{\mathrm{LWP}}+D_{\mathrm{SW}}$")
    a.set_xlabel(r"$F_{\mathrm{LW,net}}$ (W m$^{-2}$)"); a.set_ylabel(r"$D(x)$ ((W m$^{-2}$)$^2$ h$^{-1}$)")
    a.set_xlim(-100, 20); a.set_ylim(0, max(200, np.nanmax(j["D_obs"][ok]) * 1.2))
    a.legend([f1, f2, l1, l2], [h.get_label() for h in (f1, f2, l1, l2)], loc="upper right", fontsize=6)
    a.text(-0.06, 1.05, "C", transform=a.transAxes, fontsize=9, fontweight="bold")
    # D
    a = ax[1, 1]; pos = np.arange(4)
    clw = [sea[s]["share_lw"] for s in SEASONS]; csw = [sea[s]["share_sw"] for s in SEASONS]
    a.bar(pos, clw, 0.65, color=S.C_LWP, label=r"$D_{\mathrm{LWP}}$ (LW)", edgecolor="k", linewidth=0.5)
    a.bar(pos, csw, 0.65, bottom=clw, color=S.C_SW, label=r"$D_{\mathrm{SW}}$ (SW)", edgecolor="k", linewidth=0.5)
    for i, s in enumerate(SEASONS):
        lw_i, sw_i = integer_shares(clw[i], csw[i], sea[s]["r2"])
        if lw_i > 0: a.text(pos[i], clw[i] / 2, f"{lw_i}%", ha="center", va="center", fontsize=6, fontweight="bold", color="white")
        if sw_i > 0: a.text(pos[i], clw[i] + csw[i] / 2, f"{sw_i}%", ha="center", va="center", fontsize=6, fontweight="bold", color="white")
        a.text(pos[i], clw[i] + csw[i] + 2, f"R²={sea[s]['r2']:.2f}", ha="center", va="bottom", fontsize=6)
    a.set_ylabel("Variance explained (%)"); a.set_xticks(pos); a.set_xticklabels(SEASONS, fontsize=7); a.set_ylim(0, 105)
    a.legend(loc="upper center", bbox_to_anchor=(0.5, -0.12), ncol=2, fontsize=6, handlelength=1.5, columnspacing=1.0)
    a.text(-0.06, 1.05, "D", transform=a.transAxes, fontsize=9, fontweight="bold")
    fig.tight_layout(); w, h = S.finalize(fig, OUT); print(f"wrote {OUT.name} ({w:.1f}x{h:.1f}cm)")


if __name__ == "__main__":
    main()