"""Plot Fig 1 (bimodality) from ../data/fig1.pkl."""
import pickle
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import _style as S

HERE = Path(__file__).resolve().parent
D = pickle.load(open(HERE.parent / "data" / "fig1.pkl", "rb"))
OUT = HERE.parent / "figures" / "Fig1_bimodality.pdf"
CL, CD = S.CLEAR, S.CLOUDY


def main():
    fig = plt.figure(figsize=(13.6 / 2.54, 13.6 / 2.54 * 0.75))
    gt = gridspec.GridSpec(1, 2, figure=fig, left=0.08, right=0.97, top=0.97, bottom=0.68, wspace=0.30)
    gb = gridspec.GridSpec(1, 2, figure=fig, left=0.08, right=0.97, top=0.52, bottom=0.08, wspace=0.30)
    ats = [fig.add_subplot(gt[0, i]) for i in (0, 1)]
    apd = [fig.add_subplot(gb[0, i]) for i in (0, 1)]
    for ax, ts, ticks, title in ((ats[0], D["ts_djf"], ["Dec 1", "Dec 8", "Dec 15"], "DJF 2004"),
                                 (ats[1], D["ts_jja"], ["Aug 10", "Aug 17", "Aug 24"], "JJA 2019")):
        n = min(len(ts["obs"]), len(ts["era5"])); days = np.arange(n) / 24
        ax.plot(days, ts["obs"][:n], "-", lw=0.7, color=S.OBS, alpha=0.9, label="Obs")
        ax.plot(days, ts["era5"][:n], "-", lw=0.7, color=S.ERA5, alpha=0.9, label="ERA5")
        ax.axhline(-45, color=CL, ls=":", lw=1.2); ax.axhline(-5, color=CD, ls=":", lw=1.2)
        ax.set_xticks([0, 7, 14]); ax.set_xticklabels(ticks); ax.set_ylim(-90, 20)
        ax.set_title(title, fontweight="bold")
        ax.legend(loc="lower right", bbox_to_anchor=(1.0, 1.002), ncol=1, borderaxespad=0, frameon=False)
    ats[0].set_ylabel(r"$F_{\mathrm{LW,net}}$ (W m$^{-2}$)"); ats[1].tick_params(labelleft=False)
    for ax, pdf, title in ((apd[0], D["pdf_djf"], "DJF (Dec-Jan-Feb)"),
                           (apd[1], D["pdf_jja"], "JJA (Jun-Jul-Aug)")):
        for key, color, lab in (("obs", S.OBS, "Obs"), ("era5", S.ERA5, "ERA5")):
            ax.fill_between(pdf["x"], pdf[key], alpha=0.2, color=color)
            ax.plot(pdf["x"], pdf[key], "-", lw=1.5, color=color, label=lab)
        ax.axvline(-45, color=CL, ls=":", lw=1.2); ax.axvline(-5, color=CD, ls=":", lw=1.2)
        ax.text(-58, 1.95, "Clear", color=CL, fontsize=9, ha="center", va="bottom", fontweight="bold")
        ax.text(13, 1.95, "Cloudy", color=CD, fontsize=9, ha="center", va="bottom", fontweight="bold")
        ax.set_xlabel(r"$F_{\mathrm{LW,net}}$ (W m$^{-2}$)"); ax.set_xlim(-100, 30); ax.set_ylim(0, 2.8)
        ax.legend(loc="upper left"); ax.set_title(title, fontweight="bold")
    apd[0].set_ylabel(r"Probability density ($\times 10^{-2}$)"); apd[1].tick_params(labelleft=False)
    for ax, lab, y in ((ats[0], "A", 1.15), (ats[1], "B", 1.15), (apd[0], "C", 1.05), (apd[1], "D", 1.05)):
        ax.text(-0.1, y, lab, transform=ax.transAxes, fontsize=11, fontweight="bold")
    w, h = S.finalize(fig, OUT); print(f"wrote {OUT.name} ({w:.1f}x{h:.1f}cm)")


if __name__ == "__main__":
    main()