"""Figures for enso-objective (150 dpi PNG; series palette blue/orange/aqua/yellow; ink/muted/grid greys)."""
import numpy as np, pandas as pd, os, matplotlib
matplotlib.use('Agg'); import matplotlib.pyplot as plt
R = os.environ.get('ROOT', '/work/enso-objective'); B, O, A, Y = '#2a78d6', '#eb6834', '#1baf7a', '#eda100'
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})
FIG = os.environ.get('FIGDIR', f'{R}/fig'); os.makedirs(FIG, exist_ok=True)
def rv(P, y):
    cP = np.cumsum(P, axis=1)[:, :2]; cO = np.cumsum(np.eye(3)[y], axis=1)[:, :2]; return ((cP - cO) ** 2).sum(axis=1) / 2
M = pd.read_csv(f'{R}/data/matched_v5.csv')
P = lambda s, pre: s[[f'{pre}_lan', f'{pre}_neu', f'{pre}_eln']].values
# ---- fig 1: P(La Nina) and P(El Nino) for the target season 5 months after the last observed season (official k=3)
s = M[(M.match == 'same') & (M.lead == 5)].sort_values('tt').copy(); s['date'] = pd.to_datetime(dict(year=s.tt // 12, month=s.tt % 12 + 1, day=15))
fig, axs = plt.subplots(2, 1, figsize=(10, 5.6), sharex=True)
for ax, (c, j, nm) in zip(axs, [('lan', 0, 'La Niña'), ('eln', 2, 'El Niño')]):
    for dt, yy in zip(s.date, s.y):
        if yy == j: ax.axvspan(dt - pd.Timedelta(days=16), dt + pd.Timedelta(days=16), color='#eeeeee', lw=0, zorder=0)
    ax.plot(s.date, s[f'p_{c}'], color=B, lw=2, label='official CPC/IRI (human consensus)')
    ax.plot(s.date, s[f'q_{c}'], color=O, lw=2, label='IRI objective (model plume), ~1 week later')
    ax.set_ylim(0, 1); ax.set_ylabel(f'P({nm})'); ax.set_title(f'{nm}: grey = target season verified as {nm}', fontsize=9.5, loc='left', color=MUTED)
axs[0].legend(frameon=False, fontsize=8.5, loc='upper left', ncol=2)
axs[1].set_xlabel('target season (centre month); lead 5: the target season is centred 3 months after the issue month')
fig.suptitle('Same issue month, same target season: official vs objective ENSO probabilities', x=0.01, ha='left', fontsize=10.5)
fig.tight_layout(); fig.savefig(f'{FIG}/fig3_timeseries_lead5.png', dpi=150); plt.close(fig)
# ---- fig 2: RPSS vs climatology by lead, matched cells (2013-08..2025-04 official issues)
fig, ax = plt.subplots(figsize=(7.5, 4.3))
for mt, pre, col, lab, ls in (('older', 'p', B, 'official', '-'), ('same', 'q', O, 'objective, issued ~1 week after the official', '-'),
                               ('older', 'q', O, 'objective, issued ~3 weeks before the official', '--')):
    g = M[M.match == mt].groupby('lead').apply(lambda t: 1 - rv(P(t, pre), t.y.values).mean() / rv(P(t, 'c'), t.y.values).mean(), include_groups=False)
    ax.plot(g.index, g.values, 'o' + ls, color=col, lw=2, ms=5, label=lab)
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('RPSS vs 30-yr climatology')
ax.set_ylim(0, 0.85); ax.set_title('Skill on matched cells, official issues 2013-08 to 2025-04 (ONI v5)', fontsize=10, loc='left'); ax.legend(frameon=False, fontsize=8.5)
fig.tight_layout(); fig.savefig(f'{FIG}/fig1_skill_by_lead.png', dpi=150); plt.close(fig)
# ---- fig 3: mean RPS difference by issue year (objective minus official), same-month match
s = M[M.match == 'same'].copy(); s['d'] = rv(P(s, 'q'), s.y.values) - rv(P(s, 'p'), s.y.values)
g = s.groupby('iy').d.agg(['mean', 'size'])
fig, ax = plt.subplots(figsize=(8, 3.8))
ax.bar(g.index, g['mean'], color=[O if v < 0 else B for v in g['mean']], width=0.7)
ax.axhline(0, color=INK, lw=1); ax.set_xticks(g.index); ax.set_xticklabels([f'{i}' + ('*' if n < 90 else '') for i, n in zip(g.index, g['size'])], fontsize=8.5)
ax.set_ylabel('mean RPS: objective minus official'); ax.set_xlabel('issue year (* partial year)')
ax.text(0.01, 0.97, 'above zero: official (humans) better', transform=ax.transAxes, va='top', fontsize=8.5, color=B)
ax.text(0.01, 0.04, 'below zero: objective (models) better', transform=ax.transAxes, va='bottom', fontsize=8.5, color=O)
ax.set_title('Who was closer, year by year (all leads, same issue month)', fontsize=10, loc='left')
fig.tight_layout(); fig.savefig(f'{FIG}/fig2_by_year.png', dpi=150); plt.close(fig)
# ---- fig 4: temperature b by era with 95% CI
T = pd.read_csv(f'{R}/results/temper_v5.csv')
rows = [('matched-same B 2013-08..2020-10', 'official', '2013-08..2020-10\nofficial'), ('matched-same B 2013-08..2020-10', 'objective', '2013-08..2020-10\nobjective'),
        ('matched-same C 2020-11..2025-04', 'official', '2020-11..2025-04\nofficial'), ('matched-same C 2020-11..2025-04', 'objective', '2020-11..2025-04\nobjective'),
        ('objective-alone all', 'objective', '2011-12..2025-12\nobjective, all')]
fig, ax = plt.subplots(figsize=(8, 3.8))
for i, (st, fc, lab) in enumerate(rows):
    r = T[(T.set == st) & (T.forecast == fc)].iloc[0]; col = B if fc == 'official' else O
    ax.plot([r.b_lo, r.b_hi], [i, i], color=col, lw=2); ax.plot(r.b, i, 'o', color=col, ms=7)
    ax.text(r.b_hi + 0.06, i, f'{r.b:.2f}', va='center', fontsize=8.5, color=INK)
ax.axvline(1, color=INK, lw=1); ax.set_yticks(range(len(rows))); ax.set_yticklabels([r[2] for r in rows], fontsize=8.5); ax.invert_yaxis()
ax.set_xscale('log'); ax.minorticks_off(); ax.set_xticks([0.5, 1, 2, 4]); ax.set_xticklabels(['0.5', '1', '2', '4']); ax.set_xlim(0.3, 6)
ax.set_xlabel('temperature b (b > 1: should have been sharper; b < 1: over-confident)\nlines: 95% interval, issue-year bootstrap', fontsize=9)
ax.set_title('Confidence: official (blue) and objective (orange) on identical cells; objective alone below', fontsize=10, loc='left')
fig.tight_layout(); fig.savefig(f'{FIG}/fig4_confidence.png', dpi=150); plt.close(fig)
print('ok')