"""Usage: score.py [--oni v5|v6] [--ext]
Score official CPC/IRI ENSO probabilities against the ONI. Forecasts: data/official_probs_raw.csv (IRI quick-look tables,
issues 2013-08..2025-04); with --ext also CPC 'strengths' tables for issues 2025-05..2026-01 (data/cpc_strengths_lan_eln.csv;
later issues verify against the Relative ONI and are excluded). ONI: data/oni_v5.csv (ERSSTv5, default; the official index
when these forecasts were made) or data/oni_v6.csv (ERSSTv6), both tabulated to 0.1 degC.
Category of a target season: ONI <= -0.5 La Nina, >= +0.5 El Nino, else Neutral.
k = position in the 9-season table (k=0: first listed season). Target centre month = issue month + offset + k, where the
offset (0 or -1) is read from the first listed season. Baselines, using only ONI seasons observed before the issue month
(latest known season is centred two months before the issue month):
  clim30 : category frequencies of the target season over the 30 years before the issue year
  reg    : Gaussian forecast from a regression of target ONI on [latest ONI, latest minus 3 months earlier], fit for each
           (target season, lead) on 1951..(issue year - 1); residual sd from the fit.
Writes data/pairs_<tag>.csv, results/summary_<tag>.csv, results/reliability_<tag>.csv (tag = v5, v6, v6_ext)."""
import numpy as np, pandas as pd, os
from scipy.stats import norm
R = os.environ.get('ENSO_ROOT', '/work/enso-outlooks')
SEAS = ['DJF', 'JFM', 'FMA', 'MAM', 'AMJ', 'MJJ', 'JJA', 'JAS', 'ASO', 'SON', 'OND', 'NDJ']
import sys
V = sys.argv[sys.argv.index('--oni') + 1] if '--oni' in sys.argv else 'v5'; EXT = '--ext' in sys.argv; TAG = V + ('_ext' if EXT else '')
o = pd.read_csv(f'{R}/data/oni_{V}.csv')
o['cm'] = o.SEAS.map({s: i + 1 for i, s in enumerate(SEAS)}); o['t'] = o.YR * 12 + o.cm - 1  # month index of season centre
oni = o.set_index('t').ANOM
# CPC tabulates ONI to one decimal; classify on the rounded value (e.g. -0.46 -> -0.5 = La Nina)
cat = lambda x: (lambda r: np.where(r <= -0.5, 0, np.where(r >= 0.5, 2, 1)))(np.round(np.asarray(x, dtype=float) + 1e-9 * np.sign(x), 1))
d = pd.read_csv(f'{R}/data/official_probs_raw.csv')
if EXT:
    c = pd.read_csv(f'{R}/data/cpc_strengths_lan_eln.csv'); c = c[(c.issue >= '2025-05') & (c.issue <= '2026-01')].copy()
    c['k'] = c.groupby('issue').cumcount(); c['centre_month'] = c.season.map({s_: i + 1 for i, s_ in enumerate(SEAS)})
    d = pd.concat([d, c[['issue', 'k', 'season', 'centre_month', 'p_lan', 'p_neu', 'p_eln']]], ignore_index=True)
d[['p_lan', 'p_neu', 'p_eln']] = d[['p_lan', 'p_neu', 'p_eln']].fillna(0) / 100
d[['p_lan', 'p_neu', 'p_eln']] = d[['p_lan', 'p_neu', 'p_eln']].div(d[['p_lan', 'p_neu', 'p_eln']].sum(axis=1), axis=0)
d['iy'] = d.issue.str[:4].astype(int); d['im'] = d.issue.str[5:].astype(int); d['ti'] = d.iy * 12 + d.im - 1
first = d[d.k == 0].set_index('issue').centre_month
d['off'] = d.issue.map(lambda i: ((first[i] - int(i[5:]) + 6) % 12) - 6)
d['tt'] = d.ti + d.off + d.k; d['lead'] = d.tt - (d.ti - 2)  # months from latest observed season centre
d['oni'] = d.tt.map(oni); d = d.dropna(subset=['oni']).copy(); d['y'] = cat(d.oni)
# baselines
cl, rg = [], []
for _, r in d.iterrows():
    cm = r.tt % 12; hist = o[(o.cm == cm + 1) & (o.YR >= r.iy - 30) & (o.YR < r.iy)]
    c = np.bincount(cat(hist.ANOM.values), minlength=3) / len(hist); cl.append(c)
    L = int(r.lead); tr = []
    for Y in range(1951, r.iy):
        tgt = Y * 12 + cm; x0 = oni.get(tgt - L); x1 = oni.get(tgt - L - 3); yv = oni.get(tgt)
        if None not in (x0, x1, yv) and not np.isnan([x0, x1, yv]).any(): tr.append((x0, x0 - x1, yv))
    tr = np.array(tr); A = np.column_stack([np.ones(len(tr)), tr[:, 0], tr[:, 1]]); beta, *_ = np.linalg.lstsq(A, tr[:, 2], rcond=None)
    sd = np.sqrt(np.sum((tr[:, 2] - A @ beta) ** 2) / (len(tr) - 3))
    x0 = oni[r.ti - 2]; x1 = oni[r.ti - 5]; mu = beta @ [1, x0, x0 - x1]
    pl = norm.cdf(-0.5, mu, sd); pe = 1 - norm.cdf(0.5, mu, sd); rg.append([pl, 1 - pl - pe, pe])
d[['c_lan', 'c_neu', 'c_eln']] = np.array(cl); d[['r_lan', 'r_neu', 'r_eln']] = np.array(rg)
d.to_csv(f'{R}/data/pairs_{TAG}.csv', index=False)
def rps(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).mean() / 2
P = lambda s, pre: s[[f'{pre}_lan', f'{pre}_neu', f'{pre}_eln']].values
out = []
rng = np.random.default_rng(0)
for k, s in d.groupby('k'):
    y = s.y.values; o_ = dict(k=k, lead=int(s.lead.median()), n=len(s), RPS_off=rps(P(s, 'p'), y), RPS_clim=rps(P(s, 'c'), y), RPS_reg=rps(P(s, 'r'), y))
    o_['RPSS_vs_clim'] = 1 - o_['RPS_off'] / o_['RPS_clim']; o_['RPSS_vs_reg'] = 1 - o_['RPS_off'] / o_['RPS_reg']; o_['RPSS_reg_vs_clim'] = 1 - o_['RPS_reg'] / o_['RPS_clim']
    yrs = s.iy.unique(); bs = []
    for _ in range(500):
        ss = pd.concat([s[s.iy == Y] for Y in rng.choice(yrs, len(yrs))]); bs.append(1 - rps(P(ss, 'p'), ss.y.values) / rps(P(ss, 'r'), ss.y.values))
    o_['RPSS_vs_reg_lo'], o_['RPSS_vs_reg_hi'] = np.percentile(bs, [2.5, 97.5])
    for j, nm in enumerate(['lan', 'neu', 'eln']):
        o_[f'obs_{nm}'] = (y == j).mean(); o_[f'fc_{nm}'] = s[f'p_{nm}'].mean()
    out.append(o_)
S = pd.DataFrame(out); S.to_csv(f'{R}/results/summary_{TAG}.csv', index=False)
pd.set_option('display.width', 250); pd.set_option('display.max_columns', 30); print(S.round(3).to_string(index=False))
bins = [0, .05, .15, .25, .35, .45, .55, .65, .75, .85, .95, 1.0001]; rel = []
for pre, nmf in (('p', 'official'), ('r', 'regression')):
    for j, nm in enumerate(['lan', 'neu', 'eln']):
        for grp, s in (('k0-2', d[d.k <= 2]), ('k3-8', d[d.k >= 3])):
            s = s.assign(b=pd.cut(s[f'{pre}_{nm}'], bins, right=False))
            for b, h in s.groupby('b', observed=True):
                rel.append(dict(fc=nmf, cat=nm, kgrp=grp, 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_{TAG}.csv', index=False)
print(TAG, 'issues', d.issue.nunique(), d.issue.min(), d.issue.max(), 'pairs', len(d))