"""Compute the numbers behind the figures; writes results/fig*.csv."""
import sys; sys.path.insert(0, '/work/aftershock-forecasts')
from analyze import O, boot, save
import numpy as np
C = [o for o in O if o['complete']]
S = [o for o in C if o['window'] != '1 Year']
first = {}
for o in C:
    if o['eid'] not in first or int(o['creation_ms']) < first[o['eid']]: first[o['eid']] = int(o['creation_ms'])
lagmin = lambda o: (int(o['creation_ms']) - int(o['time_ms'])) / 6e4
isfirst = lambda o: int(o['creation_ms']) == first[o['eid']] and lagmin(o) <= 60   # initial forecast, issued within 1 h
F1 = [o for o in S if isfirst(o)]
print('initial forecasts (<=1 h):', len({o['eid'] for o in F1}), 'mainshocks')
edges = [0, 0.01, 0.03, 0.1, 0.3, 0.6, 1.0001]
rows = []
for name, sub in (('all versions', S), ('initial forecast', F1)):
    for lo, hi in zip(edges[:-1], edges[1:]):
        T = [o for o in sub if lo <= o['p'] < hi]
        if len(T) < 10: continue
        f, a, b = boot(T, lambda o: o['y'])
        rows.append(dict(series=name, p_lo=lo, p_hi=min(hi, 1), n=len(T), events=len({o['eid'] for o in T}),
                         meanP=np.mean([o['p'] for o in T]), obs=f, ci05=a, ci95=b))
save(rows, 'fig1_reliability.csv')
lagbin = lambda o: ('initial (<=1 h)' if isfirst(o) else '1-2 h' if o['lag_d'] < 1/12 else '2 h-1 d' if o['lag_d'] < 1 else '1-7 d' if o['lag_d'] < 7
                    else '7-30 d' if o['lag_d'] < 30 else '>30 d')
rows = []
for grp, sel in (('all', lambda o: True), ('shallow (<70 km)', lambda o: not o['deep']), ('deep (>=70 km)', lambda o: o['deep'])):
    for k in ['initial (<=1 h)', '1-2 h', '2 h-1 d', '1-7 d', '7-30 d', '>30 d']:
        T = [o for o in S if lagbin(o) == k and sel(o)]
        if not T: continue
        r, a, b = boot(T, lambda o: o['y'], lambda o: o['p'])
        rows.append(dict(group=grp, issue_lag=k, n=len(T), events=len({o['eid'] for o in T}), expected=sum(o['p'] for o in T), observed=sum(o['y'] for o in T), obs_over_exp=r, ci05=a, ci95=b))
save(rows, 'fig2_by_issue_lag.csv')
rows = []
for M in (3.0, 4.0, 5.0, 6.0):
    T = [o for o in F1 if o['M'] == M]
    r, a, b = boot(T, lambda o: o['y'], lambda o: o['p'])
    rows.append(dict(panel='magnitude', group=f'M{M:.0f}+', n=len(T), events=len({o['eid'] for o in T}), expected=sum(o['p'] for o in T), observed=sum(o['y'] for o in T), obs_over_exp=r, ci05=a, ci95=b))
for reg in ('CONUS+HI', 'Alaska', 'Marianas/Guam', 'Non-US', 'deep (>=70 km)', 'all shallow (<70 km)'):
    T = [o for o in F1 if o['deep']] if reg.startswith('deep') else [o for o in F1 if not o['deep']] if reg.startswith('all shallow') else [o for o in F1 if o['reg'] == reg]
    r, a, b = boot(T, lambda o: o['y'], lambda o: o['p'])
    rows.append(dict(panel='region', group=reg, n=len(T), events=len({o['eid'] for o in T}), expected=sum(o['p'] for o in T), observed=sum(o['y'] for o in T), obs_over_exp=r, ci05=a, ci95=b))
save(rows, 'fig3_initial_by_M_region.csv')
rows = []
for w in ['1 Day', '1 Week', '1 Month', '1 Year']:
    T = [o for o in C if o['window'] == w and np.isfinite(o['pg']) and np.isfinite(o['pb'])]
    for bname, key in (('generic model, not updated', 'pg'), ('background rate', 'pb')):
        g, a, b = boot(T, lambda o: (o['p'] - o['y']) ** 2, lambda o: (o[key] - o['y']) ** 2)
        rows.append(dict(window=w, reference=bname, n=len(T), events=len({o['eid'] for o in T}), brier_skill=1 - g, ci05=1 - b, ci95=1 - a))
save(rows, 'fig4_skill.csv')
for f in ('fig1_reliability', 'fig2_by_issue_lag', 'fig3_initial_by_M_region', 'fig4_skill'): print(open(f'results/{f}.csv').read())