"""Confidence test by temperature scaling: p_i' = p_i^b / sum_j p_j^b (one parameter for all categories).
b > 1 means the forecasts should have been sharper (under-confident); b < 1 means softer (over-confident).
Probabilities are floored at 0.005 before scaling. For each period and lead group: b fitted by minimising RPS
(full sample, issue-year block bootstrap 95% interval) and the leave-one-issue-year-out RPS change from scaling.
Usage: temper.py [tag]. Writes results/temper_<tag>.csv."""
import numpy as np, pandas as pd, os, sys
from scipy.optimize import minimize_scalar
R = os.environ.get('ENSO_ROOT', '/work/enso-outlooks'); TAG = sys.argv[1] if len(sys.argv) > 1 else 'v5'
d = pd.read_csv(f'{R}/data/pairs_{TAG}.csv'); d['period'] = np.where(d.issue < '2020-11', '2013-08..2020-10', '2020-11..')
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
def scale(P, b): Q = np.clip(P, .005, 1) ** b; return Q / Q.sum(axis=1, keepdims=True)
def fit(P, y): return minimize_scalar(lambda b: rps(scale(P, b), y), bounds=(0.2, 5), method='bounded').x
rng = np.random.default_rng(4); out = []
for pre, nm in (('p', 'official'), ('r', 'regression')):
    for per, dp in [('all', d)] + list(d.groupby('period')):
        for g, s in (('k0-2', dp[dp.k <= 2]), ('k3-5', dp[(dp.k >= 3) & (dp.k <= 5)]), ('k6-8', dp[dp.k >= 6]), ('k0-8', dp)):
            P = s[[f'{pre}_lan', f'{pre}_neu', f'{pre}_eln']].values; y = s.y.values; yr = s.iy.values; yrs = np.unique(yr)
            b = fit(P, y); bs = []
            for _ in range(200):
                idx = np.concatenate([np.where(yr == Y)[0] for Y in rng.choice(yrs, len(yrs))]); bs.append(fit(P[idx], y[idx]))
            Pcv = np.zeros_like(P)
            for Y in yrs:
                tr, te = yr != Y, yr == Y; Pcv[te] = scale(P[te], fit(P[tr], y[tr]))
            out.append(dict(forecast=nm, period=per, leads=g, n=len(s), b=b, b_lo=np.percentile(bs, 2.5), b_hi=np.percentile(bs, 97.5),
                            RPS=rps(P, y), RPS_cv_scaled=rps(Pcv, y), cv_gain=1 - rps(Pcv, y) / rps(P, y)))
o = pd.DataFrame(out); o.to_csv(f'{R}/results/temper_{TAG}.csv', index=False); pd.set_option('display.width', 200); print(o.round(3).to_string(index=False))