"""Plot Fig 6 (conditional landscapes) from ../data/fig6.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" / "fig6.pkl", "rb"))
OUT = HERE.parent / "figures" / "Fig6_conditional_landscape.pdf"
XG = D["XG"]; YE = D["Y_EVAL"]; UV = D["U_VMAX"]; XB = D["X_BARRIER"]
TS = D["T_SLICES"]; PS = D["P_SLICES"]; CL, CD = S.CLEAR, S.CLOUDY
LBL_NETLW = r"$F_{\mathrm{LW,net}}$ (W m$^{-2}$)"


def landscape(ax, res, xlabel, show_y):
    im = ax.pcolormesh(res["bin_centers"], YE, res["U"], cmap="viridis_r", shading="auto", vmin=0, vmax=UV, rasterized=True)
    bc = res["bin_centers"]
    for w, c in ((res["ww"], CD), (res["wc"], CL)):
        m = np.isfinite(w)
        if m.sum(): ax.plot(bc[m], w[m], "-o", color=c, ms=2.5, mew=0.3, lw=1.0, alpha=0.95, zorder=6)
    ax.set_xlabel(xlabel, fontsize=9); ax.tick_params(labelsize=8)
    if show_y: ax.set_ylabel(LBL_NETLW, fontsize=9)
    else: plt.setp(ax.get_yticklabels(), visible=False)
    ax.set_ylim(-90, 10); return im


def main():
    figw = 15.0 / 2.54
    fig, ax = plt.subplots(3, 2, figsize=(figw, figw * 1.16), gridspec_kw={"hspace": 0.52, "wspace": 0.30})
    im = landscape(ax[0, 0], D["resT"], r"$T_{\mathrm{2m}}$ ($^\circ$C)", True)
    landscape(ax[0, 1], D["resP"], r"$P_{\mathrm{sfc}}$ (hPa)", False)
    for a, lt in ((ax[0, 0], "A"), (ax[0, 1], "B")):
        a.text(-0.10, 1.08, lt, transform=a.transAxes, fontsize=9, fontweight="bold")
    cax = ax[0, 1].inset_axes([1.05, 0.0, 0.05, 1.0]); cb = fig.colorbar(im, cax=cax)
    cb.set_label(r"$U(x) = -\ln\, p(x)$", fontsize=7); cb.ax.tick_params(labelsize=6)
    h = [Line2D([0], [0], color=CL, marker="o", lw=1.0, ms=2.5, label="Clear"),
         Line2D([0], [0], color=CD, marker="o", lw=1.0, ms=2.5, label="Cloudy")]
    ax[0, 0].legend(handles=h, fontsize=6, loc="lower left", handletextpad=0.3, borderpad=0.3, labelspacing=0.2,
                    facecolor="k", framealpha=0.4, edgecolor="none", labelcolor="white")
    cmT = plt.cm.coolwarm(np.linspace(0.05, 0.95, len(TS))); cmP = plt.cm.viridis(np.linspace(0.05, 0.9, len(PS)))

    def row(al, ar, YT, YP, Y0, ylab, letters, zero=False):
        for c, Tv, prof in zip(cmT, TS, YT.T): al.plot(XG, prof, color=c, lw=1.0, label=rf"${Tv:.0f}\,^\circ$C")
        al.plot(XG, Y0, color="k", lw=1.7, label="all DJF")
        for c, Pv, prof in zip(cmP, PS, YP.T): ar.plot(XG, prof, color=c, lw=1.0, label=rf"${Pv:.0f}$ hPa")
        ar.plot(XG, Y0, color="k", lw=1.7, label="all DJF")
        for a, lt in ((al, letters[0]), (ar, letters[1])):
            if zero: a.axhline(0, color="k", lw=0.6)
            a.axvline(XB, color="gray", ls=":", lw=0.7); a.set_xlabel(LBL_NETLW, fontsize=9); a.tick_params(labelsize=8)
            y0, y1 = a.get_ylim()
            a.set_ylim(y0 - 0.20 * (y1 - y0), y1) if zero else a.set_ylim(y0, y1 + 0.42 * (y1 - y0))
            a.legend(fontsize=6, ncol=2, handletextpad=0.4, columnspacing=0.8, loc="upper right" if not zero else "lower left")
            a.text(-0.16 if a is al else -0.13, 1.04, lt, transform=a.transAxes, fontsize=9, fontweight="bold")
        al.set_ylabel(ylab, fontsize=9)
    row(ax[1, 0], ax[1, 1], D["a_Ts"], D["a_Ps"], D["a0"], r"$a(x)$ (W m$^{-2}$ h$^{-1}$)", "CD", zero=True)
    ax[1, 0].set_title(r"$a(F_{\mathrm{LW,net}}\,|\,T_{\mathrm{2m}})$", fontsize=9)
    ax[1, 1].set_title(r"$a(F_{\mathrm{LW,net}}\,|\,P_{\mathrm{sfc}})$", fontsize=9)
    row(ax[2, 0], ax[2, 1], D["D_Ts"], D["D_Ps"], D["D0"], r"$D(x)$ ((W m$^{-2}$)$^2$ h$^{-1}$)", "EF")
    ax[2, 0].set_title(r"$D(F_{\mathrm{LW,net}}\,|\,T_{\mathrm{2m}})$", fontsize=9)
    ax[2, 1].set_title(r"$D(F_{\mathrm{LW,net}}\,|\,P_{\mathrm{sfc}})$", fontsize=9)
    w, ht = S.finalize(fig, OUT); print(f"wrote {OUT.name} ({w:.1f}x{ht:.1f}cm)")


if __name__ == "__main__":
    main()