"""Sensitivity checks for the headline obs/exp ratios (M >= Mcat, windows <= 1 month unless stated)."""
import sys; sys.path.insert(0, '/work/aftershock-forecasts')
from analyze import O, boot
import numpy as np, collections
P = [o for o in O if o['complete'] and o['window'] != '1 Year']
def rep(name, T):
    r, a, b = boot(T, lambda o: o['y'], lambda o: o['p'])
    print(f"{name:52s} n={len(T):6d} ev={len({o['eid'] for o in T}):3d} sumP={sum(o['p'] for o in T):8.1f} sumY={sum(o['y'] for o in T):6d} obs/exp={r:.2f} [{a:.2f},{b:.2f}]")
rep('baseline', P)
# 1. one forecast per mainshock: the first version, each window separately (independent across rows within a window)
first = {}
for o in O:
    if o['complete'] and (o['eid'] not in first or int(o['creation_ms']) < first[o['eid']]): first[o['eid']] = int(o['creation_ms'])
for w in ['1 Day', '1 Week', '1 Month', '1 Year']:
    rep(f'first version only, {w}', [o for o in O if o['complete'] and o['window'] == w and int(o['creation_ms']) == first[o['eid']]])
# version issued closest to 1 day / 7 days after the mainshock
for lag in (1, 7):
    best = {}
    for o in O:
        if o['complete']:
            d = abs((int(o['creation_ms']) - int(o['time_ms'])) / 864e5 - lag)
            if o['eid'] not in best or d < best[o['eid']][0]: best[o['eid']] = (d, o['vid'])
    for w in ['1 Week', '1 Month']:
        rep(f'version nearest day {lag}, {w}', [o for o in O if o['complete'] and o['window'] == w and o['vid'] == best[o['eid']][1] and best[o['eid']][0] < 0.5])
# 2. drop the 5 / 10 mainshocks contributing most expected events
E = collections.Counter()
for o in P: E[o['eid']] += o['p']
for k in (5, 10):
    top = {e for e, _ in E.most_common(k)}
    rep(f'drop top {k} mainshocks by sum P ({",".join(sorted(top)[:3])}...)' if k == 5 else f'drop top {k} mainshocks by sum P', [o for o in P if o['eid'] not in top])
print('top 10 by sum P:', [(e, round(v, 1), sum(o['y'] for o in P if o['eid'] == e)) for e, v in E.most_common(10)])
# 3. mainshock magnitude class
for lo, hi in ((4, 5), (5, 6), (6, 10)):
    rep(f'mainshock M{lo}-{hi}', [o for o in P if lo <= float(o['mag']) < hi])
# 4. year of mainshock
for y in ('2019', '2020', '2024', '2025', '2026'):
    import time
    T = [o for o in P if time.gmtime(int(o['time_ms']) / 1000).tm_year == int(y)]
    if T: rep(f'mainshock year {y}', T)
# 5. first version only, all windows <= 1 month, by threshold magnitude and by region
F1 = [o for o in O if o['complete'] and o['window'] != '1 Year' and int(o['creation_ms']) == first[o['eid']] and o['lag_d'] <= 1 / 24]
print('--- initial forecasts = first version, issued <= 1 h after the mainshock:', len({o['eid'] for o in F1}), 'mainshocks')
for M in (3.0, 4.0, 5.0, 6.0):
    rep(f'initial, <=1mo, M>={M:.0f}', [o for o in F1 if o['M'] == M])
for reg in ('CONUS+HI', 'Alaska', 'Marianas/Guam', 'Non-US'):
    rep(f'initial, <=1mo, {reg}', [o for o in F1 if o['reg'] == reg])
    rep(f'initial, <=1mo, {reg}, shallow <70 km', [o for o in F1 if o['reg'] == reg and not o['deep']])
rep('initial, <=1mo, deep >=70 km', [o for o in F1 if o['deep']])
rep('initial, <=1mo, shallow <70 km', [o for o in F1 if not o['deep']])
rep('initial, <=1mo, shallow, excluding Marianas/Guam', [o for o in F1 if not o['deep'] and o['reg'] != 'Marianas/Guam'])
rep('all versions, <=1mo, deep >=70 km', [o for o in P if o['deep']])
rep('all versions, <=1mo, shallow, excl. Marianas/Guam', [o for o in P if not o['deep'] and o['reg'] != 'Marianas/Guam'])
for w in ['1 Day', '1 Week', '1 Month']:
    T = [o for o in O if o['complete'] and o['window'] == w and int(o['creation_ms']) == first[o['eid']] and o['lag_d'] <= 1 / 24]
    rep(f'initial (<=1 h) only, {w}', T)
for M in (4.0, 5.0):
    rep(f'initial, <=1mo, M>={M:.0f}, shallow, excl. Marianas', [o for o in F1 if o['M'] == M and not o['deep'] and o['reg'] != 'Marianas/Guam'])
lag = [(int(o['creation_ms']) - int(o['time_ms'])) / 6e4 for o in F1 if o['window'] == '1 Day' and o['M'] == float(o['Mcat_ev']) // 1 + (1 if float(o['Mcat_ev']) % 1 > 0.06 else 0)]
print('first-version issue lag (min) quantiles 10/50/90:', np.quantile(lag, [.1, .5, .9]).round(1))