"""Score the IRI objective ENSO probabilities alone (2011-12..2025-12 issues) and against the official consensus on
matched cells. Usage: score.py [--oni v5|v6]. Reads data/obj_pairs_<V>.csv, data/matched_<V>.csv.
Eras: A 2011-12..2013-07 (objective only; earlier IRI issues were forecaster judgment and are excluded), B 2013-08..2020-10 (official = averaged forecaster probabilities),
C 2020-11..2025-04 (official = calibrated from forecaster Nino3.4 means, per run 11 lit note), D 2025-05..2025-12 (objective only).
Temperature b as in run 11 (p^b renormalised, floor 0.005, RPS-minimising; b > 1 = under-confident), with issue-year
block bootstrap. Writes results/obj_skill_<V>.csv, results/matched_<V>.csv, results/temper_<V>.csv, results/reliability_<V>.csv."""
import numpy as np, pandas as pd, os, sys
from scipy.optimize import minimize_scalar
R = os.environ.get('ROOT', '/work/enso-objective')
V = sys.argv[sys.argv.index('--oni') + 1] if '--oni' in sys.argv else 'v5'
os.makedirs(f'{R}/results', exist_ok=True); pd.set_option('display.width', 250); pd.set_option('display.max_columns', 40)
def rps_v(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
rps = lambda P, y: rps_v(P, y).mean()
def scale(P, b): Q = np.clip(P, .005, 1) ** b; return Q / Q.sum(axis=1, keepdims=True)
def fitb(P, y): return minimize_scalar(lambda b: rps(scale(P, b), y), bounds=(0.2, 5), method='bounded').x
def era(i): return np.select([i < '2013-08', i < '2020-11', i < '2025-05'], ['A 2011-12..2013-07', 'B 2013-08..2020-10', 'C 2020-11..2025-04'], 'D 2025-05..2025-12')
cols = lambda pre: [f'{pre}_lan', f'{pre}_neu', f'{pre}_eln']
rng = np.random.default_rng(12)
def boot_years(s, stat, n=500):
    yrs = s.iy.unique(); g = {Y: s.index[s.iy == Y] for Y in yrs}; out = []
    for _ in range(n):
        out.append(stat(s.loc[np.concatenate([g[Y] for Y in rng.choice(yrs, len(yrs))])]))
    return np.percentile(out, [2.5, 97.5])
# ---- 1. objective alone
d = pd.read_csv(f'{R}/data/obj_pairs_{V}.csv')
# 2003-06..2011-11 issues are IRI forecaster judgment with quartile categories, not the objective Gaussian rendering (lit/literature.md 1b)
d = d[(d.issue >= '2011-12') & (d.issue <= '2025-12')].reset_index(drop=True); d['era'] = era(d.issue)
o1 = []
for (e, L), s in list(d.groupby(['era', 'lead'])) + [(('all', L), s) for L, s in d.groupby('lead')]:
    q, c = rps(s[cols('q')].values, s.y.values), rps(s[cols('c')].values, s.y.values)
    o1.append(dict(era=e, lead=L, n=len(s), RPS_obj=q, RPS_clim=c, RPSS=1 - q / c))
O1 = pd.DataFrame(o1); O1.to_csv(f'{R}/results/obj_skill_{V}.csv', index=False)
print(O1.pivot(index='lead', columns='era', values='RPSS').round(2).to_string())
# ---- 2. matched comparison
M = pd.read_csv(f'{R}/data/matched_{V}.csv'); M['era'] = era(M.issue)
M[['b_lan', 'b_neu', 'b_eln']] = (M[cols('p')].values + M[cols('q')].values) / 2  # 50/50 blend
o2 = []
for (mt, e, lg), s in [((mt, e, lg), s) for (mt, e), s0 in M.groupby(['match', 'era']) for lg, s in
                       (('2-4', s0[s0.lead <= 4]), ('5-7', s0[(s0.lead >= 5) & (s0.lead <= 7)]), ('8-10', s0[s0.lead >= 8]), ('all', s0))] + \
                      [((mt, 'B+C', lg), s) for mt, s0 in M.groupby('match') for lg, s in
                       (('2-4', s0[s0.lead <= 4]), ('5-7', s0[(s0.lead >= 5) & (s0.lead <= 7)]), ('8-10', s0[s0.lead >= 8]), ('all', s0))]:
    s = s.reset_index(drop=True); y = s.y.values
    r = dict(match=mt, era=e, leads=lg, n=len(s), RPS_off=rps(s[cols('p')].values, y), RPS_obj=rps(s[cols('q')].values, y),
             RPS_blend=rps(s[cols('b')].values, y), RPS_clim=rps(s[cols('c')].values, y))
    r['RPSS_off'] = 1 - r['RPS_off'] / r['RPS_clim']; r['RPSS_obj'] = 1 - r['RPS_obj'] / r['RPS_clim']; r['RPSS_blend'] = 1 - r['RPS_blend'] / r['RPS_clim']
    r['obj_minus_off'] = r['RPS_obj'] - r['RPS_off']
    r['diff_lo'], r['diff_hi'] = boot_years(s, lambda t: rps(t[cols('q')].values, t.y.values) - rps(t[cols('p')].values, t.y.values))
    o2.append(r)
O2 = pd.DataFrame(o2); O2.to_csv(f'{R}/results/matched_{V}.csv', index=False)
print(O2.round(3).to_string(index=False))
# ---- 3. temperature on identical cells (match = same), plus objective-alone eras
o3 = []
def tb(name, s, pre):
    s = s.reset_index(drop=True); P = s[cols(pre)].values; y = s.y.values
    lo, hi = boot_years(s, lambda t: fitb(t[cols(pre)].values, t.y.values), n=200)
    o3.append(dict(set=name, forecast=pre, n=len(s), b=fitb(P, y), b_lo=lo, b_hi=hi, RPS=rps(P, y)))
for mt in ('same', 'older'):
    for e, s in list(M[M.match == mt].groupby('era')) + [('B+C', M[M.match == mt])]:
        for pre in ('p', 'q', 'b'): tb(f'matched-{mt} {e}', s, pre)
for e, s in list(d.groupby('era')) + [('all', d)]:
    tb(f'objective-alone {e}', s, 'q')
O3 = pd.DataFrame(o3).replace({'forecast': {'p': 'official', 'q': 'objective', 'b': 'blend'}})
O3.to_csv(f'{R}/results/temper_{V}.csv', index=False); print(O3.round(3).to_string(index=False))
# ---- 4. reliability (matched-same; official vs objective) and mean-probability bias
bins = [0, .05, .15, .25, .35, .45, .55, .65, .75, .85, .95, 1.0001]; rel = []
for src, s0, pre in (('official', M[M.match == 'same'], 'p'), ('objective-matched', M[M.match == 'same'], 'q'), ('objective-all', d, 'q')):
    for e, s1 in list(s0.groupby('era')) + [('all', s0)]:
        for j, nm in enumerate(['lan', 'neu', 'eln']):
            s = s1.assign(bin=pd.cut(s1[f'{pre}_{nm}'], bins, right=False))
            for b, h in s.groupby('bin', observed=True):
                rel.append(dict(forecast=src, era=e, cat=nm, bin=str(b), n=len(h), mean_p=h[f'{pre}_{nm}'].mean(), obs=(h.y == j).mean()))
pd.DataFrame(rel).to_csv(f'{R}/results/reliability_{V}.csv', index=False)
bias = []
for src, s0, pre in (('official', M[M.match == 'same'], 'p'), ('objective-matched', M[M.match == 'same'], 'q'), ('objective-all', d, 'q')):
    for e, s in list(s0.groupby('era')) + [('all', s0)]:
        r = dict(forecast=src, era=e, n=len(s))
        for j, nm in enumerate(['lan', 'neu', 'eln']): r[f'fc_{nm}'] = s[f'{pre}_{nm}'].mean(); r[f'obs_{nm}'] = (s.y == j).mean()
        P = s[cols(pre)].values; r['sharp_maxp'] = P.max(axis=1).mean(); bias.append(r)
B = pd.DataFrame(bias); B.to_csv(f'{R}/results/bias_{V}.csv', index=False); print(B.round(3).to_string(index=False))
# ---- 5. real-time recalibrated plume mean (postproc.py) on the same cells
rc = pd.read_csv(f'{R}/data/recal_{V}.csv').set_index(['ti', 'tt'])[cols('r')].rename(columns=lambda c_: 'g' + c_[1:])
M2 = M.copy(); M2['oti'] = np.where(M2.match == 'same', M2.ti, M2.ti - 1)
M2 = M2.join(rc, on=['oti', 'tt'], how='inner')
d2 = d.join(rc, on=['ti', 'tt'], how='inner')
o5 = []
for nm, S0, pres in (('matched-same', M2[M2.match == 'same'], ('p', 'q', 'g')), ('matched-older', M2[M2.match == 'older'], ('p', 'q', 'g')),
                     ('objective-alone', d2, ('q', 'g'))):
    for e, s0 in list(S0.groupby('era')) + [('all', S0)]:
        for lg, s in (('2-4', s0[s0.lead <= 4]), ('5-7', s0[(s0.lead >= 5) & (s0.lead <= 7)]), ('8-11', s0[s0.lead >= 8]), ('all', s0)):
            s = s.reset_index(drop=True); y = s.y.values; c_ = rps(s[cols('c')].values, y)
            r = dict(set=nm, era=e, leads=lg, n=len(s))
            for pre in pres: r[f'RPSS_{pre}'] = 1 - rps(s[cols(pre)].values, y) / c_
            r['g_minus_q'] = rps(s[cols('g')].values, y) - rps(s[cols('q')].values, y)
            r['gq_lo'], r['gq_hi'] = boot_years(s, lambda t: rps(t[cols('g')].values, t.y.values) - rps(t[cols('q')].values, t.y.values), n=300)
            if 'p' in pres:
                r['g_minus_p'] = rps(s[cols('g')].values, y) - rps(s[cols('p')].values, y)
                r['gp_lo'], r['gp_hi'] = boot_years(s, lambda t: rps(t[cols('g')].values, t.y.values) - rps(t[cols('p')].values, t.y.values), n=300)
            if lg == 'all': r['b_g'] = fitb(s[cols('g')].values, y)
            o5.append(r)
O5 = pd.DataFrame(o5).rename(columns={'RPSS_p': 'RPSS_official', 'RPSS_q': 'RPSS_objective', 'RPSS_g': 'RPSS_recal'})
O5.to_csv(f'{R}/results/recal_{V}.csv', index=False); print(O5.round(3).to_string(index=False))
# ---- 6. paired difference in b (official minus objective) on identical cells, issue-year block bootstrap
o6 = []
for mt in ('same', 'older'):
    for e, s in list(M[M.match == mt].groupby('era')) + [('B+C', M[M.match == mt])]:
        s = s.reset_index(drop=True)
        f = lambda t: fitb(t[cols('p')].values, t.y.values) - fitb(t[cols('q')].values, t.y.values)
        lo, hi = boot_years(s, f, n=300); o6.append(dict(match=mt, era=e, n=len(s), b_off_minus_obj=f(s), lo=lo, hi=hi))
O6 = pd.DataFrame(o6); O6.to_csv(f'{R}/results/bdiff_{V}.csv', index=False); print(O6.round(3).to_string(index=False))