"""Reliability, 95%-range coverage, PIT, and log-score skill vs baselines (primary set: M >= Mcat, windows <= 1 month
unless stated). Baselines: generic Reasenberg-Jones (the event's own first, un-updated parameters), the 10-year
pre-mainshock background rate in the same circle, and published forecast + background combined."""
import sys; sys.path.insert(0, '/work/aftershock-forecasts')
from analyze import O, boot, save, WINS
import numpy as np, collections
P = [o for o in O if o['complete']]
S = [o for o in P if o['window'] != '1 Year']
# --- reliability: bins of stated probability
edges = [0, 0.01, 0.03, 0.1, 0.3, 0.6, 1.0001]
rel = []
for sub, name in ((S, 'complete<=1mo'), ([o for o in S if o['reg'] in ('CONUS+HI', 'Alaska')], 'complete<=1mo CONUS+AK'),
                  ([o for o in O if not o['complete'] and o['window'] != '1 Year'], 'below Mcat<=1mo')):
    for lo, hi in zip(edges[:-1], edges[1:]):
        T = [o for o in sub if lo <= o['p'] < hi]
        if len(T) < 20: continue
        f, a, b = boot(T, lambda o: o['y'])
        rel.append(dict(set=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))
for r in rel: print(f"{r['set']:24s} [{r['p_lo']:.2f},{r['p_hi']:.2f}) n={r['n']:6d} ev={r['events']:3d} P={r['meanP']:.3f} obs={r['obs']:.3f} [{r['ci05']:.3f},{r['ci95']:.3f}]")
save(rel, 'reliability.csv')
# --- 95% range coverage on informative forecasts (upper bound >= 1)
cov = []
for w in WINS:
    T = [o for o in P if o['window'] == w and o['hi'] >= 1]
    if not T: continue
    c, a, b = boot(T, lambda o: o['lo'] <= o['N'] <= o['hi'])
    cov.append(dict(window=w, n=len(T), events=len({o['eid'] for o in T}), coverage=c, ci05=a, ci95=b,
                    above=np.mean([o['N'] > o['hi'] for o in T]), below=np.mean([o['N'] < o['lo'] for o in T])))
    print(f"95% range {w:8s} n={len(T):6d} ev={cov[-1]['events']:3d} coverage={c:.3f} [{a:.3f},{b:.3f}] above={cov[-1]['above']:.3f} below={cov[-1]['below']:.3f}")
save(cov, 'coverage95.csv')
# --- PIT histogram (randomized, from the 1%..99% fractiles) for forecasts with P>=0.1
pit = np.array([o['pit'] for o in S if o['p'] >= 0.1 and np.isfinite(o['pit'])])
h = np.histogram(pit, bins=10, range=(0, 1))[0]
print('PIT deciles (P>=0.1, n=%d):' % len(pit), (h / len(pit)).round(3))
np.savetxt('results/pit_deciles.csv', np.c_[np.arange(10) / 10, h], delimiter=',', header='bin_lo,count', fmt='%.1f,%d')
# --- skill: mean log score difference per forecast (nats), event bootstrap
def ls(p, y):
    p = min(max(p, 1e-4), 1 - 1e-4); return np.log(p) if y else np.log(1 - p)
comb = lambda o: 1 - (1 - o['p']) * (1 - o['pb'])
gain = []
for sub, name in ((S, 'complete<=1mo'), ([o for o in P if o['window'] == '1 Year'], 'complete 1 Year')):
    T = [o for o in sub if np.isfinite(o['pg']) and np.isfinite(o['pb'])]
    for bname, bf in (('vs generic RJ', lambda o: o['pg']), ('vs background', lambda o: o['pb'])):
        g, a, b = boot(T, lambda o: ls(o['p'], o['y']) - ls(bf(o), o['y']))
        gain.append(dict(set=name, comparison='published ' + bname, n=len(T), events=len({o['eid'] for o in T}), gain_nats=g, ci05=a, ci95=b))
    g, a, b = boot(T, lambda o: ls(comb(o), o['y']) - ls(o['p'], o['y']))
    gain.append(dict(set=name, comparison='published+background vs published', n=len(T), events=len({o['eid'] for o in T}), gain_nats=g, ci05=a, ci95=b))
    for bname, bf in (('generic', lambda o: o['pg']), ('background', lambda o: o['pb'])):
        g, a, b = boot(T, lambda o: (o['p'] - o['y']) ** 2, lambda o: (bf(o) - o['y']) ** 2); g, a, b = 1 - g, 1 - b, 1 - a
        gain.append(dict(set=name, comparison='Brier skill vs ' + bname, n=len(T), events=len({o['eid'] for o in T}), gain_nats=g, ci05=a, ci95=b))
for r in gain: print(f"{r['set']:16s} {r['comparison']:36s} n={r['n']:6d} ev={r['events']} {r['gain_nats']:+.4f} [{r['ci05']:+.4f},{r['ci95']:+.4f}]")
save(gain, 'skill.csv')

# --- influence: 1-year skill vs background, dropping each mainshock in turn
T = [o for o in P if o['window'] == '1 Year' and np.isfinite(o['pg']) and np.isfinite(o['pb'])]
def bss(X):
    return 1 - sum((o['p'] - o['y']) ** 2 for o in X) / sum((o['pb'] - o['y']) ** 2 for o in X)
base = bss(T); infl = sorted(((bss([o for o in T if o['eid'] != e]) - base, e) for e in {o['eid'] for o in T}), reverse=True)
print(f'1-year Brier skill vs background {base:+.3f}; leave-one-mainshock-out range {base + infl[-1][0]:+.3f} to {base + infl[0][0]:+.3f}; most influential: {infl[0][1]} ({base + infl[0][0]:+.3f} without it), {infl[-1][1]} ({base + infl[-1][0]:+.3f} without it)')