"""Figures for enso-outlooks (150 dpi PNG; palette blue/orange/aqua)."""
import numpy as np, pandas as pd, os, matplotlib
matplotlib.use('Agg'); import matplotlib.pyplot as plt
R = os.environ.get('ENSO_ROOT', '/work/enso-outlooks'); B, O, A = '#2a78d6', '#eb6834', '#1baf7a'; INK, MUTED, GRID = '#222222', '#666666', '#dddddd'
plt.rcParams.update({'font.size': 10, 'axes.edgecolor': MUTED, 'xtick.color': MUTED, 'ytick.color': MUTED, 'axes.spines.top': False,
                     'axes.spines.right': False, 'axes.grid': True, 'grid.color': GRID, 'grid.linewidth': 0.6})
S = pd.read_csv(f'{R}/results/summary_v5.csv')
fig, ax = plt.subplots(figsize=(7, 4.2))
ax.plot(S.lead, S.RPSS_vs_clim, 'o-', color=B, lw=2, ms=5, label='official vs climatology')
ax.plot(S.lead, S.RPSS_reg_vs_clim, 'o-', color=O, lw=2, ms=5, label='persistence regression vs climatology')
ax.errorbar(S.lead, S.RPSS_vs_reg, yerr=[S.RPSS_vs_reg - S.RPSS_vs_reg_lo, S.RPSS_vs_reg_hi - S.RPSS_vs_reg], fmt='o-', color=A, lw=2, ms=5, elinewidth=1, label='official vs persistence regression (95% CI)')
ax.axhline(0, color=INK, lw=1); ax.set_xlabel('lead (months from the last observed ONI season to the target season)'); ax.set_ylabel('ranked probability skill score')
ax.set_title('Official CPC/IRI ENSO outlooks, 2013-08 to 2025-04 issues, verified on ONI (ERSSTv5)', fontsize=10, loc='left'); ax.legend(frameon=False, fontsize=8.5)
fig.tight_layout(); fig.savefig(f'{R}/fig/fig1_skill_by_lead.png', dpi=150); plt.close(fig)
d = pd.read_csv(f'{R}/data/pairs_v5.csv'); bins = [0, .05, .15, .25, .35, .45, .55, .65, .75, .85, .95, 1.0001]
fig, axs = plt.subplots(1, 3, figsize=(11, 3.9), sharey=True)
for ax, (j, c, nm) in zip(axs, [(0, 'lan', 'La Niña'), (1, 'neu', 'Neutral'), (2, 'eln', 'El Niño')]):
    for pre, col, lab in (('p', B, 'official'), ('r', O, 'persistence regression')):
        s = d.assign(b=pd.cut(d[f'{pre}_{c}'], bins, right=False)); g = s.groupby('b', observed=True).agg(n=('y', 'size'), p=(f'{pre}_{c}', 'mean'), o=('y', lambda y: (y == j).mean())); g = g[g.n >= 15]
        ax.plot(g.p, g.o, 'o-', color=col, lw=2, ms=5, label=lab)
    ax.plot([0, 1], [0, 1], color=INK, lw=1); ax.set_xlim(0, 1); ax.set_ylim(0, 1); ax.set_title(f'{nm}, all leads', fontsize=10, loc='left'); ax.set_xlabel('forecast probability')
axs[0].set_ylabel('observed frequency'); axs[0].legend(frameon=False, fontsize=8.5, loc='upper left')
fig.tight_layout(); fig.savefig(f'{R}/fig/fig2_reliability.png', dpi=150); plt.close(fig); print('ok')