"""Plot Fig 2 (drift / diffusion / potential) from ../data/fig2.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" / "fig2.pkl", "rb"))
OUT = HERE.parent / "figures" / "Fig2_drift_diffusion_potential.pdf"
G = D["guides"]; XB, XCL, XCD = G["x_barrier"], G["x_clear"], G["x_cloudy"]
CA, CD = S.CLEAR, S.CLOUDY
o, e = D["obs"], D["era5"]
XLIM = (-75, 20)


def guides(ax, y0, y1, line=True):
    for xv, c in ((XB, "gray"), (XCL, CA), (XCD, CD)):
        (ax.axvline(xv, color=c, ls=":", lw=1.2) if line
         else ax.plot([xv, xv], [y0, y1], ls=":", lw=1.2, color=c, clip_on=False))


def main():
    fig, ax = plt.subplots(2, 2, figsize=(13.6 / 2.54, 13.6 / 2.54 * 0.75), sharex=True)
    # A drift
    a = ax[0, 0]
    a.plot(o["x"], o["a"], "-", lw=1.0, color=S.OBS, label="Obs")
    a.plot(e["x"], e["a"], "-", lw=1.0, color=S.ERA5, label="ERA5")
    xf = np.linspace(-50, -15, 200)
    for s, col, lab in ((o, S.OBS, "Obs fit"), (e, S.ERA5, "ERA5 fit")):
        f = s["fit"]; a.plot(xf, f["slope"] * xf + f["intercept"], "-.", lw=0.8, color=col, label=lab)
    a.axhline(0, color="black", lw=0.8, alpha=0.7); guides(a, -4, 3, line=False)
    a.set_ylabel(r"$a(x)$ (W m$^{-2}$ h$^{-1}$)"); a.set_xlim(XLIM); a.set_ylim(-4, 4)
    a.legend(loc="lower left", fontsize=6)
    a.text(XCL, 3.7, "Clear", fontsize=8, color=CA, ha="center", va="top", fontweight="bold")
    a.text(XCD, 3.7, "Cloudy", fontsize=8, color=CD, ha="center", va="top", fontweight="bold")
    a.text(XB, 3.7, r"$x_0$", fontsize=8, color="gray", ha="center", va="top")
    a.text(-0.06, 1.05, "A", transform=a.transAxes, fontsize=9, fontweight="bold")
    a.text(0.5, 1.05, r"Drift $\boldsymbol{a(x)}$", transform=a.transAxes, fontsize=8, fontweight="bold", ha="center")
    # B diffusion
    a = ax[0, 1]
    a.plot(o["x"], o["D"], "-", lw=1.0, color=S.OBS, label="Obs")
    a.plot(e["x"], e["D"], "-", lw=1.0, color=S.ERA5, label="ERA5")
    guides(a, 0, 0); a.set_ylabel(r"$D(x)$ ((W m$^{-2}$)$^2$ h$^{-1}$)")
    a.set_xlim(XLIM); a.set_ylim(0, None); a.legend(loc="upper right", fontsize=6)
    a.text(-0.06, 1.05, "B", transform=a.transAxes, fontsize=9, fontweight="bold")
    a.text(0.5, 1.05, r"Diffusion $\boldsymbol{D(x)}$", transform=a.transAxes, fontsize=8, fontweight="bold", ha="center")
    # C potential components
    a = ax[1, 0]
    a.plot(o["x"], o["phi_a"], "-", lw=1.0, color=S.OBS); a.plot(e["x"], e["phi_a"], "-", lw=1.0, color=S.ERA5)
    a.plot(o["x"], o["phi_d"], "--", lw=1.0, color=S.OBS); a.plot(e["x"], e["phi_d"], "--", lw=1.0, color=S.ERA5)
    guides(a, 0, 0); a.set_xlabel(r"$F_{\mathrm{LW,net}}$ (W m$^{-2}$)"); a.set_ylabel("Potential component"); a.set_xlim(XLIM)
    h = [Line2D([0], [0], color=S.OBS, lw=1.0, label="Obs"), Line2D([0], [0], color=S.ERA5, lw=1.0, label="ERA5"),
         Line2D([0], [0], color="0.45", lw=1.0, ls="-", label=r"$\Phi_a$"), Line2D([0], [0], color="0.45", lw=1.0, ls="--", label=r"$\Phi_D$")]
    a.legend(handles=h, loc="upper center", fontsize=6, ncol=2, handlelength=1.8, columnspacing=2.0)
    a.text(-0.06, 1.05, "C", transform=a.transAxes, fontsize=9, fontweight="bold")
    a.text(0.5, 1.05, r"$\boldsymbol{\Phi_a}$ and $\boldsymbol{\Phi_D}$", transform=a.transAxes, fontsize=8, fontweight="bold", ha="center")
    # D effective potential
    a = ax[1, 1]
    a.plot(o["x"], o["u"], "-", lw=1.0, color=S.OBS, label="Obs"); a.plot(e["x"], e["u"], "-", lw=1.0, color=S.ERA5, label="ERA5")
    guides(a, 0, 0)
    icl = np.argmin(np.abs(o["x"] - XCL)); icd = np.argmin(np.abs(o["x"] - XCD))
    a.scatter([o["x"][icl], o["x"][icd]], [o["u"][icl], o["u"][icd]], c=[CA, CD], s=30, zorder=5, edgecolor="white", linewidth=0.8)
    a.set_xlabel(r"$F_{\mathrm{LW,net}}$ (W m$^{-2}$)"); a.set_ylabel(r"$U_{\mathrm{eff}}(x)$"); a.set_xlim(XLIM); a.set_ylim(0, None)
    a.legend(loc="upper center", fontsize=6)
    a.text(XCL - 3, o["u"][icl] - 0.08, "Clear", fontsize=6, color=CA, ha="right", va="top")
    a.text(XCD + 7, o["u"][icd] + 0.08, "Cloudy", fontsize=6, color=CD, ha="left", va="bottom")
    a.text(-0.06, 1.05, "D", transform=a.transAxes, fontsize=9, fontweight="bold")
    a.text(0.5, 1.05, r"$\boldsymbol{U_{\mathrm{eff}} = \Phi_a + \Phi_D}$", transform=a.transAxes, fontsize=8, fontweight="bold", ha="center")
    fig.tight_layout(); fig.subplots_adjust(hspace=0.25, wspace=0.35, bottom=0.12)
    w, h = S.finalize(fig, OUT)
    print(f"wrote {OUT.name}  ({w:.1f} x {h:.1f} cm)  tau_obs={o['fit']['tau']:.1f}h")


if __name__ == "__main__":
    main()