"""Figures for the flare-forecasts write-up (150 dpi PNG, palette blue/orange/aqua/yellow)."""
import numpy as np, pandas as pd, matplotlib
matplotlib.use('Agg'); import matplotlib.pyplot as plt
import os; R = os.environ.get('FLARE_ROOT', '/work/flare-forecasts'); B, O, A, Y = '#2a78d6', '#eb6834', '#1baf7a', '#eda100'; INK, MUTED, GRID = '#222222', '#666666', '#dddddd'
plt.rcParams.update({'font.size': 10, 'axes.edgecolor': MUTED, 'axes.labelcolor': INK, 'xtick.color': MUTED, 'ytick.color': MUTED,
                     'axes.spines.top': False, 'axes.spines.right': False, 'axes.grid': True, 'grid.color': GRID, 'grid.linewidth': 0.6})
a = pd.read_csv(f'{R}/data/flare_pairs.csv', parse_dates=['issue']); b = pd.read_csv(f'{R}/data/flare_pairs_oldscale.csv', parse_dates=['issue'])
SW = pd.Timestamp('2019-12-09')
# Fig 1: obs/exp by half-year, M and X, lead 1
fig, axs = plt.subplots(2, 1, figsize=(8, 6.2), sharex=True)
for ax, ev, minexp in ((axs[0], 'M', 8), (axs[1], 'X', 4)):
    for d, col, lab in ((a, B, 'operational GOES scale'), (b, O, 'legacy scale (GOES-16+ flux × 0.7)')):
        s = d[(d.event == ev) & (d.lead == 1)].copy(); s['h'] = s.issue.dt.year + np.where(s.issue.dt.month <= 6, 0.25, 0.75)
        g = s.groupby('h').agg(o=('y', 'sum'), e=('p', 'sum')); g = g[g.e >= minexp]
        if d is b: g = g[g.index > 2019.9]
        g = g.reindex(np.arange(1996.25, 2027, 0.5).round(2))  # NaN gaps break the line across omitted half-years
        ax.plot(g.index, g.o / g.e, 'o-', color=col, lw=2, ms=4.5, label=lab)
    ax.axhline(1, color=INK, lw=1); ax.axvline(2019.94, color=MUTED, lw=1, ls='--')
    ax.set_ylabel(f'{ev}-class: observed / forecast'); ax.set_ylim(0, 1.6 if ev == 'M' else 1.2)
    if ev == 'M': ax.text(2019.94 + 0.2, 0.08, 'GOES-16 becomes\nprimary (2019-12-09)', color=MUTED, fontsize=8.5)
axs[0].legend(loc='upper left', frameon=False, fontsize=9)
axs[0].set_title('Event days per forecast event day, day-1 SWPC probabilities, half-years\n(1 = calibrated; <1 = over-forecast; half-years with few expected events omitted)', fontsize=10, loc='left')
fig.tight_layout(); fig.savefig(f'{R}/fig/fig1_obs_exp_by_halfyear.png', dpi=150); plt.close(fig)
# Fig 2: matched activity
sb = pd.read_csv(f'{R}/results/scale_bins.csv'); sb = sb[(sb.event == 'M') & (sb.n >= 40)]
mid = lambda s: s.str.extract(r'\((\d+), (\d+)\]').astype(float).mean(axis=1)
sb['x'] = mid(sb.fbin)
fig, axs = plt.subplots(1, 2, figsize=(9.5, 4.2), sharey=True)
ax = axs[0]
for per, col, lab in (('C24', B, 'Cycle 24 (legacy GOES)'), ('C25a', O, '2019-12 to 2024, operational scale')):
    s = sb[sb.period == per]; ax.plot(s.x, s.freq_oper, 'o-', color=col, lw=2, ms=5, label=lab)
s = sb[sb.period == 'C25a']; ax.plot(s.x, s.freq_old, 'o--', color=O, lw=1.5, ms=5, mfc='white', label='2019-12 to 2024, legacy scale')
ax.set_title('Observed: fraction of days with an M flare', fontsize=10, loc='left'); ax.set_xlabel('F10.7 90-day mean (sfu)'); ax.set_ylabel('probability'); ax.legend(frameon=False, fontsize=8.5, loc='upper left')
ax = axs[1]
for per, col, lab in (('C24', B, 'Cycle 24'), ('C25a', O, '2019-12 to 2024'), ('C25b', A, '2025 to 2026-10')):
    s = sb[sb.period == per]; ax.plot(s.x, s.mean_p, 'o-', color=col, lw=2, ms=5, label=lab)
ax.set_title('Forecast: mean SWPC day-1 M probability', fontsize=10, loc='left'); ax.set_xlabel('F10.7 90-day mean (sfu)'); ax.legend(frameon=False, fontsize=8.5, loc='upper left')
fig.tight_layout(); fig.savefig(f'{R}/fig/fig2_matched_activity.png', dpi=150); plt.close(fig)
# Fig 3: reliability, M and X, lead 1, before/after switch
bins = [0, .015, .075, .125, .225, .325, .425, .525, .625, .725, 1.0]
fig, axs = plt.subplots(1, 2, figsize=(9.5, 4.4))
for ax, ev in zip(axs, 'MX'):
    for per, lo, hi, col, lab in (('pre', '1996-01-01', SW, B, '1999 to 2019-12-08'), ('c25a', SW, '2025-01-01', O, '2019-12-09 to 2024'), ('c25b', '2025-01-01', '2027-01-01', A, '2025 to 2026-10')):
        s = a[(a.event == ev) & (a.lead == 1) & (a.issue >= lo) & (a.issue < hi)].copy(); s['bin'] = pd.cut(s.p, bins, include_lowest=True)
        g = s.groupby('bin', observed=True).agg(n=('y', 'size'), p=('p', 'mean'), o=('y', 'mean')); g = g[g.n >= 25]
        se = np.sqrt(g.o * (1 - g.o) / g.n)
        ax.errorbar(g.p, g.o, yerr=1.96 * se, fmt='o-', color=col, lw=2, ms=5, capsize=0, elinewidth=1, label=lab)
    lim = 1.0 if ev == 'M' else 0.6
    ax.plot([0, lim], [0, lim], color=INK, lw=1); ax.set_xlim(0, lim); ax.set_ylim(0, lim)
    ax.set_xlabel(f'forecast probability ({ev}-class, day 1)'); ax.set_ylabel('observed frequency')
    ax.set_title(f'{ev}-class reliability (bins with ≥25 forecasts; ±1.96 SE)', fontsize=10, loc='left')
axs[0].legend(frameon=False, fontsize=8.5, loc='upper left')
fig.tight_layout(); fig.savefig(f'{R}/fig/fig3_reliability.png', dpi=150); plt.close(fig)
print('ok')