"""Build forecast/outcome pairs for the IRI objective (model-based) ENSO probabilities and match them to the official
CPC/IRI consensus scored in run 11 (/work/enso-outlooks/data/pairs_<V>.csv).
Usage: pairs.py [--oni v5|v6|roni]  (roni = CPC Relative ONI, data/oni_roni.csv, rounded to 0.1 before thresholds)
Objective issue month M (title 'Mid-<Mon> YYYY'): first season (k=0) is centred on M+1, so target centre tt = M + 1 + k.
Lead is counted as in run 11: months from the latest observed season centre (M-2) to the target centre, so objective
k=0 has lead 3. Category: ONI rounded to 0.1, <= -0.5 La Nina, >= +0.5 El Nino (run 11's rule).
Baseline clim30: target-season category frequencies over the 30 years before the issue year (run 11's rule).
Writes data/obj_pairs_<V>.csv (all objective issues with a verifying ONI) and data/matched_<V>.csv with two matchings:
  same  : official issue M vs objective issue M (objective issued ~1 week later, same target)
  older : official issue M vs objective issue M-1 (objective issued ~3 weeks earlier, same target)."""
import numpy as np, pandas as pd, os, sys
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']
MON = 'Jan Feb Mar Apr May Jun Jul Aug Sep Oct Nov Dec'.split()
o = pd.read_csv(f'{R}/data/oni_roni.csv' if V == 'roni' else f'{E}/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
oni = o.set_index('t').ANOM
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/objective_probs_raw.csv').rename(columns={'#0000ff': 'q_lan', '#aaaaaa': 'q_neu', '#ff0000': 'q_eln'})
d = d[['issue', 'title', 'k', 'season', 'q_lan', 'q_neu', 'q_eln']].copy()
# as-issued corrections (lit/verify_asissued.md): the server holds the wrong forecast for 2018-01/02 and revised 2013-10 values
cor = pd.read_csv(f'{R}/data/asissued_corrections.csv').set_index(['issue', 'season'])
for (i_, s_), r_ in cor.iterrows():
    m_ = (d.issue == i_) & (d.season == s_); assert m_.sum() == 1, (i_, s_)
    d.loc[m_, ['q_lan', 'q_neu', 'q_eln']] = [r_.lan, r_.neu, r_.eln]
print('as-issued corrections applied:', len(cor), 'cells')
# issue month from the title (file name month is the URL month + 1 and should agree)
d['iy'] = d.title.str[-4:].astype(int); d['im'] = d.title.str[4:7].map({m: i + 1 for i, m in enumerate(MON)})
ok = d.issue == d.iy.astype(str) + '-' + d.im.map('{:02d}'.format)  # future URL months return the latest issue: drop
print('dropped (title month != URL month):', sorted(d.issue[~ok].unique())); d = d[ok].copy()
d['ti'] = d.iy * 12 + d.im - 1; d['tt'] = d.ti + 1 + d.k
assert (d.tt % 12 == d.season.map({s: i for i, s in enumerate(SEAS)})).all(), 'season label does not match M+1+k'
d['lead'] = d.tt - (d.ti - 2)
P = d[['q_lan', 'q_neu', 'q_eln']].values / 100; d[['q_lan', 'q_neu', 'q_eln']] = P / P.sum(axis=1, keepdims=True)
d['psum'] = P.sum(axis=1) * 100
d = d[d.psum > 0].copy()  # one empty cell (2004-09 JJA 2005)
d['oni'] = d.tt.map(oni); d = d.dropna(subset=['oni']).copy(); d['y'] = cat(d.oni)
cl = []
for _, r in d.iterrows():
    hist = o[(o.cm == r.tt % 12 + 1) & (o.YR >= r.iy - 30) & (o.YR < r.iy)]
    cl.append(np.bincount(cat(hist.ANOM.values), minlength=3) / len(hist))
d[['c_lan', 'c_neu', 'c_eln']] = np.array(cl)
d.to_csv(f'{R}/data/obj_pairs_{V}.csv', index=False)
print('objective pairs', len(d), 'issues', d.issue.nunique(), d.issue.min(), d.issue.max(), 'psum range', d.psum.min(), d.psum.max())
f = pd.read_csv(f'{E}/data/pairs_v5.csv')  # official forecasts (run 11); outcome and climatology recomputed for index V
f['oni'] = f.tt.map(oni); f = f.dropna(subset=['oni']).copy(); f['y'] = cat(f.oni)
cl = []
for _, r in f.iterrows():
    hist = o[(o.cm == r.tt % 12 + 1) & (o.YR >= r.iy - 30) & (o.YR < r.iy)]
    cl.append(np.bincount(cat(hist.ANOM.values), minlength=3) / len(hist))
f[['c_lan', 'c_neu', 'c_eln']] = np.array(cl)
q = d.set_index(['ti', 'tt'])[['q_lan', 'q_neu', 'q_eln', 'issue']].rename(columns={'issue': 'obj_issue'})
out = []
for nm, shift in (('same', 0), ('older', 1)):
    m = f.copy(); m['oti'] = m.ti - shift
    m = m.join(q, on=['oti', 'tt'], how='inner'); m['match'] = nm; out.append(m)
M = pd.concat(out, ignore_index=True)
M.to_csv(f'{R}/data/matched_{V}.csv', index=False)
for nm, s in M.groupby('match'):
    print(nm, 'matched pairs', len(s), 'issues', s.issue.nunique(), 'leads', sorted(s.lead.unique()))