"""Figures for the Cycle 25 forecast-calibration audit. Run with /work/venv-plot/bin/python.
Reads scored.csv (score_karak.py), ../c24/c24_table2.csv, ../baselines/{results_c25,leadtime_curve}.csv."""
import csv, os, math
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, 'fig')
BLUE, ORANGE, AQUA, YELLOW = '#2a78d6', '#eb6834', '#1baf7a', '#eda100'
INK, MUTED, GRID = '#1f2328', '#57606a', '#d0d7de'
OBS25, OBS24 = 160.9, 81.9
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, 'axes.axisbelow': True})

S = list(csv.DictReader(open(os.path.join(HERE, 'scored.csv'), encoding='utf-8')))
def f(x):
    try: return float(x)
    except ValueError: return math.nan
for s in S:
    s['v'] = f(s['value_v2']); s['ep'] = f(s['err_plus']); s['em'] = f(s['err_minus']); s['y'] = int(s['year'])

# ---------- Fig 1: forecasts vs publication year ----------
fig, ax = plt.subplots(figsize=(9, 5.2))
rng = np.random.default_rng(1)
ax.axhspan(115.7, 236.8, color=GRID, alpha=0.35, lw=0)
ax.text(2006.2, 232, 'climatology 68% range (cycles 1-24)', color=MUTED, fontsize=8, va='top')
for s in S:
    x = s['y'] + rng.uniform(0.1, 0.9)
    if math.isnan(s['ep']):
        ax.plot(x, s['v'], 'o', ms=4, mfc='white', mec=MUTED, mew=1, zorder=3)
        continue
    hit = s['v'] - s['em'] <= OBS25 <= s['v'] + s['ep']
    c = BLUE if hit else ORANGE
    ax.errorbar(x, s['v'], yerr=[[s['em']], [s['ep']]], fmt='o', ms=4, color=c, mec='white', mew=0.6, elinewidth=1.2, capsize=0, zorder=3)
# time-matched early-rise baseline (68% interval) at issue times
L = [r for r in csv.DictReader(open(os.path.join(HERE, '..', 'baselines', 'leadtime_curve.csv')))
     if r['fit'] == '10-24' and r['predictor'] == 'ms' and r['log'] == 'False']
xs = [int(r['issue'][:4]) + (int(r['issue'][5:7]) - 0.5) / 12 for r in L]
ax.fill_between(xs, [f(r['lo68']) for r in L], [f(r['hi68']) for r in L], color=AQUA, alpha=0.25, lw=0, step=None)
ax.plot(xs, [f(r['point']) for r in L], color=AQUA, lw=2)
ax.text(xs[-1] + 0.15, f(L[-1]['point']), 'early-rise\nbaseline (68%)', color=INK, fontsize=8, va='center')
ax.axhline(OBS25, color=INK, lw=1.2, ls='--')
ax.text(2009.6, OBS25 + 3, 'observed 160.9 (Oct 2024)', color=INK, fontsize=8.5)
ax.axvline(2019.96, color=MUTED, lw=0.8, ls=':')
ax.text(2020.05, 25, 'minimum\nDec 2019', color=MUTED, fontsize=8)
ax.set_xlim(2006, 2026.6); ax.set_ylim(20, 260)
ax.set_xlabel('publication year (forecasts jittered within the year)')
ax.set_ylabel('predicted Cycle 25 maximum (SN v2, 13-month smoothed)')
from matplotlib.lines import Line2D
h = [Line2D([], [], marker='o', color=BLUE, ls='', ms=5, label='stated ±error contains 160.9'),
     Line2D([], [], marker='o', color=ORANGE, ls='', ms=5, label='stated ±error misses 160.9'),
     Line2D([], [], marker='o', mfc='white', mec=MUTED, ls='', ms=5, label='no error stated')]
ax.legend(handles=h, loc='lower right', frameon=False, fontsize=8.5)
ax.set_title('127 published Cycle 25 forecasts (Karak 2026 Table 1) against the outcome', color=INK, fontsize=11, loc='left')
fig.tight_layout(); fig.savefig(os.path.join(OUT, 'fig1_forecasts_vs_year.png'), dpi=150); plt.close(fig)

# ---------- Fig 2: coverage vs error-bar multiplier, C25 and C24 ----------
from math import erf
class norm:
    @staticmethod
    def cdf(x): return 0.5 * (1 + np.vectorize(erf)(np.asarray(x) / np.sqrt(2)))
def zs(rows, obs, key_v, key_ep, key_em):
    out = []
    for r in rows:
        v, ep, em = r[key_v], r[key_ep], r[key_em]
        if math.isnan(ep): continue
        out.append(abs(obs - v) / (ep if obs > v else em))
    return np.array(out)
z25 = zs(S, OBS25, 'v', 'ep', 'em')
C24 = []
for r in csv.DictReader(open(os.path.join(HERE, '..', 'c24', 'c24_table2.csv'))):
    e = f(r['err']); C24.append(dict(v=float(r['value']), ep=e, em=e))
z24 = zs(C24, OBS24, 'v', 'ep', 'em')
k = np.linspace(0, 4, 401)
fig, ax = plt.subplots(figsize=(6.4, 4.6))
ax.plot(k, 2 * norm.cdf(k) - 1, color=MUTED, lw=1.5, ls='--')
ax.text(2.35, 0.955, 'calibrated (Gaussian, ±err = 1σ)', color=MUTED, fontsize=8.5, va='top')
ax.plot(k, [(z25 <= kk).mean() for kk in k], color=ORANGE, lw=2)
ax.plot(k, [(z24 <= kk).mean() for kk in k], color=BLUE, lw=2)
ax.text(3.05, (z25 <= 3).mean() - 0.07, f'Cycle 25 (n={len(z25)})', color=INK, fontsize=9)
ax.text(3.3, (z24 <= 3).mean() - 0.1, f'Cycle 24 (n={len(z24)})', color=INK, fontsize=9)
for kk in (1, 2):
    ax.axvline(kk, color=GRID, lw=1)
ax.set_xlabel('error bar multiplier k (interval = forecast ± k × stated error)')
ax.set_ylabel('fraction of intervals containing the outcome')
ax.set_xlim(0, 4); ax.set_ylim(0, 1.0)
ax.set_title('Published error bars are about 2-3x too narrow', color=INK, fontsize=11, loc='left')
fig.tight_layout(); fig.savefig(os.path.join(OUT, 'fig2_coverage_curve.png'), dpi=150); plt.close(fig)

# ---------- Fig 3: stated error vs publication year, against honest baseline error ----------
fig, ax = plt.subplots(figsize=(8, 4.6))
E = [s for s in S if not math.isnan(s['ep'])]
ax.plot([s['y'] + rng.uniform(0.1, 0.9) for s in E], [(s['ep'] + s['em']) / 2 for s in E], 'o', ms=4.5, color=BLUE, mec='white', mew=0.6)
ax.axhline(np.median([(s['ep'] + s['em']) / 2 for s in E]), color=BLUE, lw=1.2, ls='--')
ax.text(2006.2, np.median([(s['ep'] + s['em']) / 2 for s in E]) + 2, 'median stated error 17', color=INK, fontsize=8.5)
ax.plot(xs, [f(r['rmse_loo']) for r in L], color=ORANGE, lw=2, drawstyle='steps-post')
ax.hlines(57.8, 2006, 2019.96, color=ORANGE, lw=2)
ax.text(2006.2, 60.5, 'honest error of climatology (sd of cycles 1-24 maxima, 58)', color=INK, fontsize=8.5)
ax.text(xs[3] + 0.2, f(L[3]['rmse_loo']) + 3, 'leave-one-cycle-out RMSE of\nearly-rise baseline at issue time', color=INK, fontsize=8.5)
ax.axvline(2019.96, color=MUTED, lw=0.8, ls=':')
ax.set_xlim(2006, 2025.5); ax.set_ylim(0, 75)
ax.set_xlabel('publication year'); ax.set_ylabel('error (SN v2)')
ax.set_title('Stated errors did not shrink as information arrived', color=INK, fontsize=11, loc='left')
fig.tight_layout(); fig.savefig(os.path.join(OUT, 'fig3_stated_error_vs_time.png'), dpi=150); plt.close(fig)
print('ok')

# ---------- Fig 4: CRPS of published forecasts vs baselines ----------
fig, ax = plt.subplots(figsize=(8, 4.6))
ax.plot([s['y'] + rng.uniform(0.1, 0.9) for s in S], [f(s['crps']) for s in S], 'o', ms=4.5, color=BLUE, mec='white', mew=0.6)
ax.hlines(16.2, 2006, 2026, color=ORANGE, lw=2)
ax.text(2006.2, 17.5, 'climatology 178.7 ± 61 (CRPS 16.2)', color=INK, fontsize=8.5)
ax.hlines(14.4, 2014.3, 2026, color=YELLOW, lw=2)
ax.text(2014.4, 11.0, 'even→odd regression (14.4)', color=INK, fontsize=8.5)
ax.plot([2022.5, 2023.5, 2024.5], [16.6, 14.0, 10.7], 's-', color=AQUA, lw=2, ms=6)
ax.text(2022.3, 4.5, 'early-rise baseline\nk = 24, 36, 48 months', color=INK, fontsize=8.5)
ax.set_xlim(2006, 2026); ax.set_ylim(0, 110)
ax.set_xlabel('publication year'); ax.set_ylabel('CRPS (SN v2; lower is better)')
ax.set_title('Most published forecasts scored worse than an average-cycle forecast', color=INK, fontsize=11, loc='left')
fig.tight_layout(); fig.savefig(os.path.join(OUT, 'fig4_crps.png'), dpi=150); plt.close(fig)
print('fig4 ok')