"""Score RSGA middle-latitude geomagnetic category probabilities (Active K4 / Minor K5 / Major-severe K6+,
remainder = quiet-unsettled K0-3) for days 1-3. PROXY truth: planetary Kp from GFZ (daily max, rounded to integer),
not Fredericksburg K which SWPC verifies against (not retrievable here).
Metrics: ranked probability score (4 ordered categories) and Brier score for P(storm = K>=5).
Baselines known at issue day D: clim365 = category frequencies over the 365 days ending D;
cond = frequencies of the target-day category conditional on day D's category, from all earlier years (expanding)."""
import numpy as np, pandas as pd
import os; R = os.environ.get('FLARE_ROOT', '/work/flare-forecasts')
g = pd.read_csv(f'{R}/raw/Kp_ap_Ap_SN_F107_since_1932.txt', comment='#', sep=r'\s+', header=None)
kp = g[list(range(7, 15))].max(axis=1)
cat = pd.Series(np.clip(np.round(kp).astype(int) - 3, 0, 3).values, index=pd.to_datetime(dict(year=g[0], month=g[1], day=g[2])))  # 0 q,1 act,2 min,3 maj
oh = pd.get_dummies(cat).reindex(columns=[0, 1, 2, 3], fill_value=0).astype(float)
clim = oh.rolling(365, min_periods=300).mean()
rs = pd.read_csv(f'{R}/data/rsga.csv', parse_dates=['issue_date', 'day1'])
rows = []
for _, r in rs.iterrows():
    D = r.issue_date
    if pd.isna(r.day1) or (r.day1 - D).days != 1 or D not in cat.index: continue
    for lead in (1, 2, 3):
        T = D + pd.Timedelta(days=lead)
        if T not in cat.index: continue
        a, mi, mj = r[f'mid_act{lead}'], r[f'mid_min{lead}'], r[f'mid_maj{lead}']
        if pd.isna(a) or pd.isna(mi) or pd.isna(mj): continue
        p = np.array([100 - a - mi - mj, a, mi, mj]) / 100
        rows.append(dict(issue=D, lead=lead, y=int(cat[T]), today=int(cat[D]), p0=p[0], p1=p[1], p2=p[2], p3=p[3],
                         c0=clim.at[D, 0], c1=clim.at[D, 1], c2=clim.at[D, 2], c3=clim.at[D, 3]))
d = pd.DataFrame(rows).dropna(); d['year'] = d.issue.dt.year
# conditional-on-today baseline, expanding window by year
for k in range(4): d[f'k{k}'] = np.nan
for lead in (1, 2, 3):
    s = d[d.lead == lead]
    for yr in sorted(s.year.unique()):
        tr = s[s.year < yr]
        if len(tr) < 365: continue
        tab = pd.crosstab(tr.today, tr.y).reindex(index=range(4), columns=range(4), fill_value=0) + 0.5
        tab = tab.div(tab.sum(axis=1), axis=0)
        te = s[s.year == yr]
        d.loc[te.index, [f'k{k}' for k in range(4)]] = tab.loc[te.today].values
d = d.dropna()
def rps(P, y):
    cP = np.cumsum(P, axis=1); cO = np.cumsum(np.eye(4)[y], axis=1); return ((cP - cO) ** 2).sum(axis=1)[:, :3].mean() if False else ((cP - cO) ** 2)[:, :3].sum(axis=1).mean() / 3
def cols(pre): return [f'{pre}{k}' for k in range(4)]
out = []
for lead, s in d.groupby('lead'):
    y = s.y.values; o = dict(lead=lead, n=len(s), freq_q=(y == 0).mean(), freq_act=(y == 1).mean(), freq_min=(y == 2).mean(), freq_maj=(y == 3).mean())
    o['mean_p_act'], o['mean_p_min'], o['mean_p_maj'] = s.p1.mean(), s.p2.mean(), s.p3.mean()
    R0 = rps(s[cols('p')].values, y)
    for b, pre in (('clim365', 'c'), ('cond', 'k')):
        o[f'RPSS_vs_{b}'] = 1 - R0 / rps(s[cols(pre)].values, y)
    st = (y >= 2).astype(float); ps = s.p2 + s.p3
    o['storm_obs_exp'] = st.sum() / ps.sum()
    o['BSS_storm_vs_clim365'] = 1 - np.mean((ps - st) ** 2) / np.mean((s.c2 + s.c3 - st) ** 2)
    o['BSS_storm_vs_cond'] = 1 - np.mean((ps - st) ** 2) / np.mean((s.k2 + s.k3 - st) ** 2)
    out.append(o)
o = pd.DataFrame(out); o.to_csv(f'{R}/results/geomag_summary.csv', index=False)
pd.set_option('display.width', 250); pd.set_option('display.max_columns', 30); print(o.round(3).to_string(index=False))
yr = []
for y_, s in d[d.lead == 1].groupby('year'):
    y = s.y.values; st = (y >= 2).astype(float); ps = s.p2 + s.p3
    yr.append(dict(year=y_, n=len(s), RPSS_clim=1 - rps(s[cols('p')].values, y) / rps(s[cols('c')].values, y),
                   RPSS_cond=1 - rps(s[cols('p')].values, y) / rps(s[cols('k')].values, y), storm_days=int(st.sum()), storm_exp=ps.sum()))
yr = pd.DataFrame(yr); yr.to_csv(f'{R}/results/geomag_by_year.csv', index=False); print(yr.round(2).to_string(index=False))
bins = [0, .015, .075, .125, .225, .325, .425, .525, 1.0]
s = d.assign(ps=d.p2 + d.p3, st=(d.y >= 2).astype(float))
rel = s.groupby(['lead', pd.cut(s.ps, bins, include_lowest=True)], observed=True).agg(n=('st', 'size'), mean_p=('ps', 'mean'), obs=('st', 'mean')).reset_index()
rel['ps'] = rel.ps.astype(str); rel.to_csv(f'{R}/results/geomag_storm_reliability.csv', index=False); print(rel[rel.lead == 1].round(3).to_string(index=False))