"""Calibration and skill of USGS operational aftershock forecasts. Reads outcomes.csv; writes results/*.csv and prints tables.
Primary set: bins with M >= the event's stated catalog completeness Mcat (USGS's own), windows 1 Day/1 Week/1 Month/1 Year.
All confidence intervals: bootstrap over mainshocks (2000 resamples), because versions of one sequence are strongly dependent."""
import sys; sys.path.insert(0, '/work/aftershock-forecasts')
import csv, os, numpy as np, collections
from region import regions
os.chdir('/work/aftershock-forecasts'); os.makedirs('results', exist_ok=True)
REG = regions()
O = [o for o in csv.DictReader(open('outcomes.csv'))]
DEPTH = {v['eid']: float(v['depth']) for v in csv.DictReader(open('versions.csv'))}
for o in O:
    o['depth'] = DEPTH[o['eid']]; o['deep'] = o['depth'] >= 70
    o['lag_d'] = (int(o['creation_ms']) - int(o['time_ms'])) / 864e5   # issue time after the mainshock, days
    o['reg'] = REG[o['eid']]; o['M'] = float(o['M']); o['p'] = float(o['prob']); o['N'] = int(o['N'])
    o['y'] = int(o['N'] >= 1); o['pg'] = float(o['p_generic']); o['pb'] = float(o['p_bg'])
    o['complete'] = o['Mcat_ev'] != '' and o['M'] >= float(o['Mcat_ev']) - 0.05
    o['hi'] = float(o['hi95']); o['lo'] = float(o['lo95']); o['pit'] = float(o['pit']); o['d0'] = float(o['d0'])
WINS = ['1 Day', '1 Week', '1 Month', '1 Year']
rng = np.random.default_rng(1)
def boot(S, num, den=None, B=2000):
    """sum(num)/sum(den) over rows S (den defaults to 1 per row, i.e. a mean), with a 90% interval from resampling
    mainshocks with replacement. num, den: functions of a row."""
    ids = {}; e = np.array([ids.setdefault(o['eid'], len(ids)) for o in S])
    n = np.bincount(e, [num(o) for o in S], len(ids)); d = np.bincount(e, [den(o) for o in S] if den else None, len(ids))
    W = rng.multinomial(len(ids), np.ones(len(ids)) / len(ids), B)
    with np.errstate(invalid='ignore', divide='ignore'): vals = (W @ n) / (W @ d)
    vals = vals[np.isfinite(vals)]
    return n.sum() / d.sum(), np.quantile(vals, 0.05), np.quantile(vals, 0.95)
def table_calibration(S, label, keyf, keys):
    rows = []
    for k in keys:
        T = [o for o in S if keyf(o) == k]
        if len(T) < 20: continue
        r, lo, hi = boot(T, lambda o: o['y'], lambda o: o['p'])
        rows.append(dict(set=label, group=k, n=len(T), events=len({o['eid'] for o in T}), meanP=np.mean([o['p'] for o in T]),
                         obs=np.mean([o['y'] for o in T]), obs_over_exp=r, ci05=lo, ci95=hi))
    return rows
def show(rows):
    for r in rows:
        print(f"{r['set']:28s} {str(r['group']):22s} n={r['n']:6d} ev={r['events']:3d} P={r['meanP']:.3f} obs={r['obs']:.3f} obs/exp={r['obs_over_exp']:.2f} [{r['ci05']:.2f},{r['ci95']:.2f}]")
def save(rows, name):
    with open(f'results/{name}', 'w', newline='') as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0])); w.writeheader()
        for r in rows: w.writerow({k: (round(v, 4) if isinstance(v, float) else v) for k, v in r.items()})
if __name__ == '__main__':
    P = [o for o in O if o['complete']]
    print('primary set:', len(P), 'forecasts,', len({o['eid'] for o in P}), 'mainshocks; all bins:', len(O))
    rows = []
    rows += table_calibration(P, 'complete, by window', lambda o: o['window'], WINS)
    rows += table_calibration([o for o in P if o['window'] != '1 Year'], 'complete<=1mo, by region', lambda o: o['reg'], ['CONUS+HI', 'Alaska', 'Marianas/Guam', 'Puerto Rico', 'Non-US'])
    rows += table_calibration([o for o in P if o['window'] != '1 Year'], 'complete<=1mo, by M', lambda o: o['M'], [3.0, 4.0, 5.0, 6.0, 7.0])
    rows += table_calibration([o for o in O if not o['complete'] and o['window'] != '1 Year'], 'BELOW Mcat<=1mo, by region', lambda o: o['reg'], ['CONUS+HI', 'Alaska', 'Marianas/Guam', 'Puerto Rico', 'Non-US'])
    leadbin = lambda o: '<2h' if o['lag_d'] < 1/12 else '2h-1d' if o['lag_d'] < 1 else '1-7d' if o['lag_d'] < 7 else '7-30d' if o['lag_d'] < 30 else '>30d'
    rows += table_calibration([o for o in P if o['window'] != '1 Year'], 'complete<=1mo, by issue lag', leadbin, ['<2h', '2h-1d', '1-7d', '7-30d', '>30d'])
    rows += table_calibration([o for o in P if o['window'] != '1 Year'], 'complete<=1mo, by model', lambda o: o['model'].split('(')[-1].rstrip(')') if 'Reas' in o['model'] else 'ETAS', ['Generic', 'Bayesian Combination', 'Sequence Specific', 'ETAS'])
    rows += table_calibration([o for o in P if o['window'] != '1 Year'], 'complete<=1mo, by depth', lambda o: 'deep >=70 km' if o['deep'] else 'shallow <70 km', ['shallow <70 km', 'deep >=70 km'])
    rows += table_calibration([o for o in P if o['window'] != '1 Year' and not o['deep']], 'complete<=1mo shallow, by region', lambda o: o['reg'], ['CONUS+HI', 'Alaska', 'Marianas/Guam', 'Puerto Rico', 'Non-US'])
    show(rows); save(rows, 'calibration_tables.csv')