"""Plot Fig 8 panels a-c (CMIP6 comparison) from ../data/fig8.pkl.
The full Fig 8 (with schematic panel D) is the vector scheme.pdf."""
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" / "fig8.pkl", "rb"))
OUT = HERE.parent / "figures" / "Fig8_cmip6_panels_abc.pdf"
DATASETS = {"obs_3hr": ("Observations", "black", "-", 1.5), "gfdl_cm4": ("GFDL-CM4", "#9467bd", "-", 1.0),
            "canesm5": ("CanESM5", "#d62728", "-", 1.0), "ipsl_cm6a_lr": ("IPSL-CM6A-LR", "#1f77b4", "-", 1.0),
            "miroc6": ("MIROC6", "#2ca02c", "-", 1.0), "cnrm_cm6_1": ("CNRM-CM6-1", "#e377c2", "-", 1.0),
            "ec_earth3": ("EC-Earth3", "#17becf", "--", 1.0)}


def main():
    x = D["x_grid"]
    fig, ax = plt.subplots(1, 3, figsize=(13.6 / 2.54, 13.6 / 2.54 * 0.36))
    panels = (("pdf", "Probability density", "A"), ("drift", r"Drift $a(x)$ (W m$^{-2}$ h$^{-1}$)", "B"),
              ("diffusion", r"$D(x)$ ((W m$^{-2}$)$^2$ h$^{-1}$)", "C"))
    for a, (field, ylab, lab) in zip(ax, panels):
        for key, (name, color, ls, lw) in DATASETS.items():
            y = D[f"{key}_{field}"]
            if not np.all(np.isnan(y)):
                a.plot(x, y, color=color, ls=ls, lw=lw, label=name)
        if field == "drift": a.axhline(0, color="gray", ls="-", lw=0.5)
        if field == "diffusion": a.set_ylim(bottom=0)
        a.set_xlabel(r"$F_{\mathrm{LW,net}}$ (W m$^{-2}$)"); a.set_ylabel(ylab, fontsize=7); a.set_xlim(-90, 10)
        a.text(-0.06, 1.10, lab, transform=a.transAxes, fontsize=9, fontweight="bold")
    h, l = ax[0].get_legend_handles_labels()
    fig.legend(h, l, loc="lower center", ncol=4, fontsize=6, bbox_to_anchor=(0.5, -0.12), frameon=False)
    fig.tight_layout(); w, ht = S.finalize(fig, OUT); print(f"wrote {OUT.name} ({w:.1f}x{ht:.1f}cm)")


if __name__ == "__main__":
    main()