"""Usage: score_flares.py [--oldscale]
Score RSGA M/X flare probabilities (day 1-3) against GOES XRA events in SWPC daily event lists.
Outcome M = at least one flare >= M1.0 (M or X) on the target UT day; X = at least one >= X1.0.
Baselines (all use only data up to and including the issue day D; target = D+lead):
  clim365 : fraction of event days in the 365 days ending D
  rate27  : fraction of event days in the 27 days ending D (one solar rotation)
  logit   : logistic regression on [event today, frac last 3 d, frac last 27 d, frac last 365 d],
            refit each calendar year on all earlier pairs (expanding window).
Writes data/flare_pairs.csv and results/flare_*.csv; prints a summary."""
import numpy as np, pandas as pd, os, sys
from sklearn.linear_model import LogisticRegression
import os; R = os.environ.get('FLARE_ROOT', '/work/flare-forecasts'); os.makedirs(f'{R}/results', exist_ok=True)
rs = pd.read_csv(f'{R}/data/rsga.csv', parse_dates=['issue_date', 'day1'])
fd = pd.read_csv(f'{R}/data/flare_days.csv', parse_dates=['date']).set_index('date').sort_index()
# --oldscale: classify flares on the legacy GOES scale (GOES-16+ fluxes x 0.7), i.e. post-2019-12 'M1' needs M1.43 operational.
old = '--oldscale' in sys.argv; col = 'peak_flux_old' if old else 'peak_flux'
ev = pd.DataFrame(index=pd.date_range(fd.index.min(), fd.index.max()))
ev['have'] = fd.reindex(ev.index)['n_xra'].notna()
ev['M'] = (fd.reindex(ev.index)[col] >= 1e-5).astype(float).where(ev['have'])
ev['X'] = (fd.reindex(ev.index)[col] >= 1e-4).astype(float).where(ev['have'])
feat = {}
for e in 'MX':
    s = ev[e]
    feat[e] = pd.DataFrame({'today': s, 'f3': s.rolling(3, min_periods=2).mean(), 'f27': s.rolling(27, min_periods=20).mean(),
                            'f365': s.rolling(365, min_periods=300).mean()})
rows = []
for _, r in rs.iterrows():
    D = r.issue_date
    if pd.isna(r.day1) or (r.day1 - D).days != 1: continue
    for lead in (1, 2, 3):
        T = D + pd.Timedelta(days=lead)
        if T not in ev.index or not ev.at[T, 'have']: continue
        for e in 'MX':
            p = r[f'{e}{lead}']
            if pd.isna(p) or D not in feat[e].index: continue
            f = feat[e].loc[D]
            rows.append(dict(issue=D, target=T, lead=lead, event=e, p=p / 100, y=ev.at[T, e], today=f.today, f3=f.f3, f27=f.f27, f365=f.f365))
pr = pd.DataFrame(rows).dropna(subset=['y', 'today', 'f3', 'f27', 'f365'])
pr['clim365'] = pr.f365; pr['rate27'] = pr.f27
pr['year'] = pr.issue.dt.year; pr['logit'] = np.nan
X = lambda d: np.column_stack([d.today, d.f3, d.f27, d.f365])
for (e, lead), g in pr.groupby(['event', 'lead']):
    for yr in sorted(g.year.unique()):
        tr = g[g.year < yr]; te = g[g.year == yr]
        if len(tr) < 365 or tr.y.nunique() < 2: continue
        m = LogisticRegression(C=1e3, max_iter=2000).fit(X(tr), tr.y)
        pr.loc[te.index, 'logit'] = m.predict_proba(X(te))[:, 1]
tag = '_oldscale' if old else ''
pr.to_csv(f'{R}/data/flare_pairs{tag}.csv', index=False)
ev_ = pr.dropna(subset=['logit'])
def bs(p, y): return float(np.mean((p - y) ** 2))
def skill(a, b): return 1 - a / b if b > 0 else np.nan
out = []
for (e, lead), g in ev_.groupby(['event', 'lead']):
    d = dict(event=e, lead=lead, n=len(g), base_rate=g.y.mean(), mean_p=g.p.mean(), obs_over_exp=g.y.sum() / g.p.sum(), BS_swpc=bs(g.p, g.y))
    for b in ('clim365', 'rate27', 'logit'):
        d[f'BS_{b}'] = bs(g[b], g.y); d[f'BSS_vs_{b}'] = 1 - d['BS_swpc'] / d[f'BS_{b}']
    # block bootstrap by calendar month for BSS vs logit
    g = g.assign(mo=g.issue.dt.to_period('M')); mos = g.mo.unique(); rng = np.random.default_rng(0)
    gm = {m: x for m, x in g.groupby('mo')}; bss = []
    for _ in range(300):
        s = pd.concat([gm[m] for m in rng.choice(mos, len(mos))])
        bss.append(1 - bs(s.p, s.y) / bs(s.logit, s.y))
    d['BSS_vs_logit_lo'], d['BSS_vs_logit_hi'] = np.percentile(bss, [2.5, 97.5])
    out.append(d)
summ = pd.DataFrame(out); summ.to_csv(f'{R}/results/flare_summary{tag}.csv', index=False)
pd.set_option('display.width', 250); pd.set_option('display.max_columns', 30)
print(summ.round(3).to_string(index=False))
# reliability (lead 1)
bins = [0, .015, .075, .125, .225, .325, .425, .525, .625, .725, .825, 1.0]
rel = []
for (e, lead), g in ev_.groupby(['event', 'lead']):
    g = g.assign(bin=pd.cut(g.p, bins, include_lowest=True))
    for b, h in g.groupby('bin', observed=True):
        rel.append(dict(event=e, lead=lead, bin=str(b), n=len(h), mean_p=h.p.mean(), obs=h.y.mean()))
rel = pd.DataFrame(rel); rel.to_csv(f'{R}/results/flare_reliability{tag}.csv', index=False)
print(rel[rel.lead == 1].round(3).to_string(index=False))
# by year, lead 1
yr = []
for (e, y), g in ev_[ev_.lead == 1].groupby(['event', 'year']):
    yr.append(dict(event=e, year=y, n=len(g), base=g.y.mean(), mean_p=g.p.mean(), obs_exp=g.y.sum() / max(g.p.sum(), 1e-9),
                   BSS_clim365=skill(bs(g.p, g.y), bs(g.clim365, g.y)), BSS_rate27=skill(bs(g.p, g.y), bs(g.rate27, g.y)), BSS_logit=skill(bs(g.p, g.y), bs(g.logit, g.y))))
yr = pd.DataFrame(yr); yr.to_csv(f'{R}/results/flare_by_year{tag}.csv', index=False)
print(yr.round(2).to_string(index=False))