"""Plot Fig 4 (LWP closure) from ../data/fig4.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" / "fig4.pkl", "rb"))
OUT = HERE.parent / "figures" / "Fig4_lwp_closure.pdf"
v = D["valid"]; CT, CS = S.C_LWP, S.C_SW
FMAX_LW, K_LWP = 83.8, 0.0446


def tprime(lwp):
    return FMAX_LW * K_LWP * np.exp(-K_LWP * lwp)


def main():
    fig = plt.figure(figsize=(13.6 / 2.54, 13.6 / 2.54 * 0.75))
    gs = fig.add_gridspec(2, 5, hspace=0.65, wspace=0.4, width_ratios=[1, 1, 0.7, 1, 1])
    aa = fig.add_subplot(gs[0, :2]); ab = fig.add_subplot(gs[0, 3:]); ac = fig.add_subplot(gs[1, 1:4])
    # A
    ax, ax2 = aa, aa.twinx(); lg = np.linspace(1, 60, 100)
    ax.plot(lg, tprime(lg), "-", lw=1.0, color=CT, label=rf"$S_{{\mathrm{{LW}}}}$ (R$^2$ = {D['r2_slw']:.2f})")
    ax2.plot(lg, D["sigma_coeff"] * np.sqrt(lg), "-", lw=1.0, color=CS, label=rf"$\sigma_{{\mathrm{{LWP}}}}$ (R$^2$ = {D['r2_sigma']:.2f})")
    ax.scatter(D["lwp_mean"][v], D["t_prime"][v], c=CT, s=20, edgecolor="white", linewidth=0.5, zorder=5)
    ax2.scatter(D["lwp_mean"][v], D["lwp_std"][v], c=CS, s=20, edgecolor="white", linewidth=0.5, zorder=5)
    ax.set_xlabel(r"LWP (g m$^{-2}$)"); ax.set_ylabel(r"$S_{\mathrm{LW}}$ (W m$^{-2}$ / g m$^{-2}$)", color=CT)
    ax.tick_params(axis="y", labelcolor=CT); ax2.set_ylabel(r"$\sigma_{\mathrm{LWP}}$ (g m$^{-2}$)", color=CS); ax2.tick_params(axis="y", labelcolor=CS)
    ax.set_xlim(0, 60); ax.set_ylim(0, 6); ax2.set_ylim(0, 70); ax.set_xticks([0, 20, 40, 60]); ax.set_yticks([0, 2, 4, 6]); ax2.set_yticks([0, 20, 40, 60])
    ax.text(-0.06, 1.10, "A", transform=ax.transAxes, fontsize=9, fontweight="bold")
    ax.legend(loc="upper left", fontsize=6); ax2.legend(loc="upper left", bbox_to_anchor=(0, 0.9), fontsize=6)
    # B
    ax, ax2 = ab, ab.twinx(); vl = np.isfinite(D["t_prime"])
    ax.scatter(D["x"][v], D["t_prime"][v], c=CT, s=20, edgecolor="white", linewidth=0.5, alpha=0.5, zorder=4)
    ax2.scatter(D["x"][v], D["lwp_std"][v], c=CS, s=20, edgecolor="white", linewidth=0.5, alpha=0.5, zorder=4)
    ax.plot(D["x"][vl], D["t_prime"][vl], "-", lw=1.0, color=CT, zorder=5, label=r"$S_{\mathrm{LW}}$")
    ax2.plot(D["x"][vl], D["lwp_std_smooth"][vl], "-", lw=1.0, color=CS, zorder=5, label=r"smoothed $\sigma_{\mathrm{LWP}}$")
    ax.set_xlabel(r"$F_{\mathrm{LW,net}}$ (W m$^{-2}$)"); ax.set_ylabel(r"$S_{\mathrm{LW}}$ (W m$^{-2}$ / g m$^{-2}$)", color=CT)
    ax.tick_params(axis="y", labelcolor=CT); ax2.set_ylabel(r"$\sigma_{\mathrm{LWP}}$ (g m$^{-2}$)", color=CS); ax2.tick_params(axis="y", labelcolor=CS)
    ax.set_xlim(-75, 10); ax.set_xticks([-60, -40, -20, 0]); ax.set_yticks([0, 1, 2, 3]); ax2.set_yticks([0, 20, 40, 60])
    ax.text(-0.06, 1.10, "B", transform=ax.transAxes, fontsize=9, fontweight="bold")
    ax.legend(loc="center left", bbox_to_anchor=(0, 0.6), fontsize=6); ax2.legend(loc="center left", bbox_to_anchor=(0, 0.5), fontsize=6)
    # C
    ax = ac
    ax.plot(D["x"][v], D["D_binned"][v], "k-", lw=1.0, label=r"$D_{\mathrm{obs}}$")
    ax.plot(D["x"][v], D["D_physics"][v], "-", lw=1.0, color=CT, label=r"$D_{\mathrm{LWP}}$")
    ax.set_xlabel(r"$F_{\mathrm{LW,net}}$ (W m$^{-2}$)"); ax.set_ylabel(r"$D(x)$ ((W m$^{-2}$)$^2$ h$^{-1}$)")
    ax.set_xlim(-75, 10); ax.set_ylim(0, max(np.nanmax(D["D_binned"][v]), np.nanmax(D["D_physics"][v])) * 1.1)
    ax.legend(loc="upper left", fontsize=6)
    ax.text(0.95, 0.95, f"r = {D['r']:.2f}\nR² = {D['r']**2:.2f}", transform=ax.transAxes, fontsize=7, va="top", ha="right")
    ax.text(-0.02, 1.10, "C", transform=ax.transAxes, fontsize=9, fontweight="bold")
    w, h = S.finalize(fig, OUT); print(f"wrote {OUT.name} ({w:.1f}x{h:.1f}cm)")


if __name__ == "__main__":
    main()