"""Calibration summaries: per category and lead group, mean forecast vs observed frequency, and the logistic
recalibration slope b in logit(P(obs)) = a + b*logit(p) (b < 1 over-confident, b > 1 under-confident), with
issue-year block-bootstrap 95% intervals, for all issues and split at 2020-11. Usage: calib.py [tag]. Writes results/calibration_<tag>.csv."""
import numpy as np, pandas as pd, os
from sklearn.linear_model import LogisticRegression
R = os.environ.get('ENSO_ROOT', '/work/enso-outlooks')
import sys
TAG = sys.argv[1] if len(sys.argv) > 1 else 'v5'
d = pd.read_csv(f'{R}/data/pairs_{TAG}.csv'); rng = np.random.default_rng(3)
# periods: before / from 2020-11 (CPC's strength table, linked to the calibrated method of L'Heureux et al. 2019)
d['period'] = np.where(d.issue < '2020-11', '2013-08..2020-10', '2020-11..')
def slope(p, y):
    x = np.log(np.clip(p, .01, .99) / (1 - np.clip(p, .01, .99))).reshape(-1, 1)
    return LogisticRegression(C=1e4, max_iter=2000).fit(x, y).coef_[0][0] if 0 < y.sum() < len(y) else np.nan
out = []
for pre, nm in (('p', 'official'), ('r', 'regression')):
    groups = [(f'{per} {g}', s) 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))]
    for g, s in groups:
        for j, c in enumerate(['lan', 'neu', 'eln']):
            p = s[f'{pre}_{c}'].values; y = (s.y.values == j).astype(int)
            yrs = s.iy.unique(); bs = []
            for _ in range(300):
                idx = np.concatenate([np.where(s.iy.values == Y)[0] for Y in rng.choice(yrs, len(yrs))]); bs.append(slope(p[idx], y[idx]))
            out.append(dict(forecast=nm, leads=g, cat=c, n=len(s), mean_p=p.mean(), obs=y.mean(), slope=slope(p, y), slope_lo=np.nanpercentile(bs, 2.5), slope_hi=np.nanpercentile(bs, 97.5)))
o = pd.DataFrame(out); o.to_csv(f'{R}/results/calibration_{TAG}.csv', index=False); pd.set_option('display.width', 200); print(o.round(2).to_string(index=False))