"""Sensitivity of the matched-activity odds ratios (scale_test.py) to the activity covariate and the reference period
(after the fresh-context review, which first ran this). Lead 1. For each specification: OR (2019-12-09..2024 vs reference)
for operational outcome, legacy-scale outcome and SWPC probability, and the gaps oper/SWPC and legacy/SWPC.
Also: month-block bootstrap of the change in legacy-scale observed/forecast (2019-12-09..2024 vs mid-1997..2019-12-08).
Writes results/scale_sensitivity.csv and results/legacy_oe_shift.csv."""
import numpy as np, pandas as pd
from sklearn.linear_model import LogisticRegression
import os; R = os.environ.get('FLARE_ROOT', '/work/flare-forecasts')
rs = pd.read_csv(f'{R}/data/rsga.csv', parse_dates=['issue_date'])[['issue_date', 'f107_90d', 'f107_obs']].rename(columns={'issue_date': 'issue'})
g = pd.read_csv(f'{R}/raw/Kp_ap_Ap_SN_F107_since_1932.txt', comment='#', sep=r'\s+', header=None)
g = pd.DataFrame({'issue': pd.to_datetime(dict(year=g[0], month=g[1], day=g[2])), 'SN': g[24].where(g[24] >= 0)})
g['SN27'] = g.SN.rolling(27, min_periods=20).mean(); g['SN81'] = g.SN.rolling(81, min_periods=60).mean()
a = pd.read_csv(f'{R}/data/flare_pairs.csv', parse_dates=['issue']); b = pd.read_csv(f'{R}/data/flare_pairs_oldscale.csv', parse_dates=['issue'])
E = ['1996-01-01', '2009-01-01', '2019-12-09', '2025-01-01', '2027-01-01']; L = ['C23', 'C24', 'C25a', 'C25b']
E2 = ['1996-01-01', '2002-01-01', '2009-01-01', '2014-05-01', '2019-12-09', '2025-01-01', '2027-01-01']; L2 = ['C23rise', 'C23dec', 'C24rise', 'C24dec', 'C25a', 'C25b']
SPECS = [('F10.7 90-day mean (main)', E, L, 'C24', ['f107_90d']), ('daily F10.7', E, L, 'C24', ['f107_obs']),
         ('F10.7 90-day + daily', E, L, 'C24', ['f107_90d', 'f107_obs']), ('sunspot number 81-day mean', E, L, 'C24', ['SN81']),
         ('sunspot number 27-day + F10.7 90-day', E, L, 'C24', ['SN27', 'f107_90d']),
         ('F10.7 90-day; reference = Cycle 24 rise (2009 to 2014-04)', E2, L2, 'C24rise', ['f107_90d']),
         ('F10.7 90-day; reference = Cycle 24 decline (2014-05 to 2019-12-08)', E2, L2, 'C24dec', ['f107_90d'])]
rows = []
for ev in 'MX':
    d = a[(a.event == ev) & (a.lead == 1)][['issue', 'p', 'y']].merge(b[(b.event == ev) & (b.lead == 1)][['issue', 'y']].rename(columns={'y': 'y_old'}), on='issue').merge(rs, on='issue').merge(g, on='issue').dropna()
    for name, edges, labels, ref, cov in SPECS:
        dd = d.assign(period=pd.cut(d.issue, pd.to_datetime(edges), labels=labels, right=False))
        Xd = pd.get_dummies(dd.period).drop(columns=ref).astype(float)
        for c in cov: Xd[c] = np.log(dd[c] + 1)
        j = list(Xd.columns).index('C25a'); r = dict(event=ev, spec=name)
        for nm, y in (('oper', dd.y), ('legacy', dd.y_old)):
            r[f'OR_{nm}'] = np.exp(LogisticRegression(C=1e4, max_iter=5000).fit(Xd.values, y.values).coef_[0][j])
        XX = np.vstack([Xd.values] * 2); yy = np.r_[np.ones(len(dd)), np.zeros(len(dd))]; ww = np.r_[dd.p.values, 1 - dd.p.values]
        r['OR_swpc'] = np.exp(LogisticRegression(C=1e4, max_iter=5000).fit(XX, yy, sample_weight=ww).coef_[0][j])
        r['gap_oper'] = r['OR_oper'] / r['OR_swpc']; r['gap_legacy'] = r['OR_legacy'] / r['OR_swpc']; rows.append(r)
s = pd.DataFrame(rows); s.to_csv(f'{R}/results/scale_sensitivity.csv', index=False)
pd.set_option('display.width', 200); print(s.round(2).to_string(index=False))
# legacy-scale obs/exp shift, bootstrap by month
d = b[(b.event == 'M') & (b.lead == 1)].copy(); d['post'] = d.issue >= '2019-12-09'; d = d[d.issue < '2025-01-01']
mo = d.issue.dt.to_period('M'); grp = {m: np.where(mo == m)[0] for m in mo.unique()}
pre_k = [m for m in grp if m < pd.Period('2019-12', 'M')]; post_k = [m for m in grp if m >= pd.Period('2019-12', 'M')]
y, p, post = d.y.values, d.p.values, d.post.values
def shift(idx_pre, idx_post): return (y[idx_post].sum() / p[idx_post].sum()) / (y[idx_pre].sum() / p[idx_pre].sum())
ip = np.where(~post)[0]; iq = np.where(post)[0]; pt = shift(ip, iq); rng = np.random.default_rng(2); bs = []
for _ in range(1000):
    i1 = np.concatenate([grp[pre_k[k]] for k in rng.integers(0, len(pre_k), len(pre_k))]); i1 = i1[~post[i1]]
    i2 = np.concatenate([grp[post_k[k]] for k in rng.integers(0, len(post_k), len(post_k))]); i2 = i2[post[i2]]
    bs.append(shift(i1, i2))
o = pd.DataFrame([dict(oe_legacy_pre=y[ip].sum() / p[ip].sum(), oe_legacy_post=y[iq].sum() / p[iq].sum(), ratio=pt, lo=np.percentile(bs, 2.5), hi=np.percentile(bs, 97.5))])
o.to_csv(f'{R}/results/legacy_oe_shift.csv', index=False); print(o.round(3).to_string(index=False))