"""Real-time recalibration of the plume mean (a forecast anyone could have made at the time from public data).
For each objective issue month M and lead, fit on issues whose target season had been observed before M
(target centre <= M-2, i.e. strictly past verification): ONI = a + b * COMBINED_AVG + e, with e ~ N(0, s^2) by lead
(pooled over target seasons; min 60 training pairs, so the first forecasts are from ~2007).
Probabilities: Gaussian CDF at -0.5/+0.5 applied to the ONI scale (category rule of run 11 on rounded ONI ignored here;
thresholds at -0.45/+0.45 would mimic rounding: we use +-0.45 to match 'rounded to 0.1, <= -0.5').
Usage: postproc.py [--oni v5|v6]. Writes data/recal_<V>.csv (issue, k, lead, tt, mu_raw, mu, s, r_lan, r_neu, r_eln)."""
import numpy as np, pandas as pd, os, sys
from scipy.stats import norm
R = os.environ.get('ROOT', '/work/enso-objective'); E = os.environ.get('ENSO_OUTLOOKS', '/work/enso-outlooks')  # run-11 ONI tables and official pairs
V = sys.argv[sys.argv.index('--oni') + 1] if '--oni' in sys.argv else 'v5'
SEAS = ['DJF', 'JFM', 'FMA', 'MAM', 'AMJ', 'MJJ', 'JJA', 'JAS', 'ASO', 'SON', 'OND', 'NDJ']
o = pd.read_csv(f'{R}/data/oni_roni.csv' if V == 'roni' else f'{E}/data/oni_{V}.csv'); o['t'] = o.YR * 12 + o.SEAS.map({s: i for i, s in enumerate(SEAS)})
oni = o.set_index('t').ANOM
pl = pd.read_csv(f'{R}/data/plume.csv'); pl = pl[pl.model == 'COMBINED AVG'].copy()
MON = 'Jan Feb Mar Apr May Jun Jul Aug Sep Oct Nov Dec'.split()
pl['iy'] = pl.title.str[-4:].astype(int); pl['im'] = pl.title.str[:3].map({m: i + 1 for i, m in enumerate(MON)})
pl = pl[pl.issue == pl.iy.astype(str) + '-' + pl.im.map('{:02d}'.format)].copy()  # drop future-month duplicates
pl['ti'] = pl.iy * 12 + pl.im - 1
c = pl.season.map({s: i for i, s in enumerate(SEAS)})
pl['dt'] = (c - (pl.ti % 12)) % 12; pl['tt'] = pl.ti + pl.dt; pl['lead'] = pl.tt - (pl.ti - 2)
pl['oni'] = pl.tt.map(oni)
out = []
for _, r in pl.iterrows():
    tr = pl[(pl.lead == r.lead) & (pl.tt <= r.ti - 2) & pl.oni.notna()]
    if len(tr) < 60: continue
    A = np.column_stack([np.ones(len(tr)), tr['nino34']]); beta, *_ = np.linalg.lstsq(A, tr.oni, rcond=None)
    s = np.sqrt(((tr.oni - A @ beta) ** 2).sum() / (len(tr) - 2)); mu = beta[0] + beta[1] * r.nino34
    pl_, pe = norm.cdf(-0.45, mu, s), 1 - norm.cdf(0.45, mu, s)
    out.append(dict(issue=r.issue, ti=r.ti, tt=r.tt, lead=r.lead, mu_raw=r.nino34, mu=mu, s=s, a=beta[0], b=beta[1], ntrain=len(tr),
                    r_lan=pl_, r_neu=1 - pl_ - pe, r_eln=pe))
D = pd.DataFrame(out); D.to_csv(f'{R}/data/recal_{V}.csv', index=False)
print(len(D), 'recalibrated cells', D.issue.min(), D.issue.max())
print(D.groupby('lead')[['a', 'b', 's']].last().round(2).to_string())