"""Plot Fig 3 (noise-induced bifurcation) from ../data/fig3.pkl."""
import pickle
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import _style as S

HERE = Path(__file__).resolve().parent
D = pickle.load(open(HERE.parent / "data" / "fig3.pkl", "rb"))
OUT = HERE.parent / "figures" / "Fig3_noise_bifurcation.pdf"
al = D["alphas"]; ao = D["alpha_obs"]; CL, CD = S.CLEAR, S.CLOUDY


def main():
    fig, ax = plt.subplots(1, 2, figsize=(13.6 / 2.54, 13.6 / 2.54 * 0.45))
    a = ax[0]
    fc, fy = D["flat"][:, 0], D["flat"][:, 1]; pc, py = D["interp"][:, 0], D["interp"][:, 1]
    a.plot(al, fy, "--", color="#888888", lw=1, alpha=0.5); a.plot(al, fc, "--", color="#888888", lw=1, alpha=0.5)
    bi = np.isfinite(pc) & np.isfinite(py)
    a.plot(al[~bi], py[~bi], "-o", color="gray", lw=1.0, ms=3)
    a.plot(al[bi], py[bi], "-o", color=CD, lw=1.0, ms=3); a.plot(al[bi], pc[bi], "-o", color=CL, lw=1.0, ms=3)
    if np.isfinite(D["obs_clear"]):
        for y, c in ((D["obs_clear"], CL), (D["obs_cloudy"], CD)):
            a.axhline(y, color=c, ls=":", lw=1, alpha=0.7)
            a.scatter([ao], [y], marker="s", s=40, c=c, edgecolor="black", linewidth=1, zorder=10)
    if np.isfinite(D["era5_peak"]):
        a.scatter([1.0], [D["era5_peak"]], marker="s", s=40, c="gray", edgecolor="black", linewidth=1, zorder=10)
    a.plot([1, 1], [-70, 0], color=S.ERA5, lw=1.0); a.text(1, 6, "ERA5", color=S.ERA5, fontsize=7, ha="center", va="bottom")
    a.plot([ao, ao], [-70, 0], color="black", lw=1.0); a.text(ao, 6, "Obs", color="black", fontsize=7, ha="center", va="bottom")
    a.set_xlabel(r"Deformation parameter $\alpha$"); a.set_ylabel(r"$F_{\mathrm{LW,net}}$ (W m$^{-2}$)")
    a.set_xlim(0.3, 2.5); a.set_ylim(-70, 10)
    h = [Line2D([0], [0], color=CL, marker="o", lw=1.0, ms=3, label="Clear"),
         Line2D([0], [0], color=CD, marker="o", lw=1.0, ms=3, label="Cloudy"),
         Line2D([0], [0], color="gray", marker="o", lw=1.0, ms=3, label="Unimodal"),
         Line2D([0], [0], color="#888888", ls="--", lw=1.0, label="Flat $D$ (control)")]
    a.legend(handles=h, loc="lower left", fontsize=6, ncol=1, handlelength=1.5)
    a.text(-0.06, 1.05, "A", transform=a.transAxes, fontsize=9, fontweight="bold")
    a = ax[1]; xf = D["x_full"]
    a.fill_between(xf, D["p_obs"], alpha=0.12, color="black")
    a.plot(xf, D["p_obs"], "-", lw=1.5, color="black", label="Obs (data)")
    a.plot(xf, D["p_sde"], "--", lw=1.0, color="#9467bd", label="SDE")
    a.set_xlabel(r"$F_{\mathrm{LW,net}}$ (W m$^{-2}$)"); a.set_ylabel(r"Probability density ($\times 10^{-2}$)")
    a.set_xlim(-100, 30); a.set_ylim(0, None); a.legend(loc="upper left", fontsize=6)
    a.text(0.5, 0.05, f"$r$ = {D['r']:.2f}\n$R^2$ = {D['r']**2:.2f}", transform=a.transAxes, fontsize=6,
           ha="center", va="bottom", bbox=dict(boxstyle="round", facecolor="white", alpha=0.8, edgecolor="gray"))
    a.text(-0.06, 1.05, "B", 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()