"""Check own Reasenberg-Jones implementation (outcomes.rj_prob logic) against USGS's stated probabilities for 300 random versions."""
import json, glob, random, numpy as np, os
os.chdir('/work/aftershock-forecasts')
Z = np.linspace(-5, 5, 41); WZ = np.exp(-Z ** 2 / 2); WZ /= WZ.sum()
def rj_prob(a, b, p, c, asig, Mm, M, t0, t1):
    I = ((t1 + c) ** (1 - p) - (t0 + c) ** (1 - p)) / (1 - p) if abs(p - 1) > 1e-9 else np.log((t1 + c) / (t0 + c))
    return float(1 - (WZ * np.exp(-10 ** (a + asig * Z + b * (Mm - M)) * I)).sum())
random.seed(1); res = {}
for f in random.sample(sorted(glob.glob('raw/fc/*.json')), 300):
    try: d = json.load(open(f))
    except Exception: continue
    m = d['model']
    if 'Reasenberg' not in m['name']: continue
    P = m['parameters']; eid = os.path.basename(f).rsplit('_', 2)[0]
    tm = json.load(open(f'raw/events/{eid}.json'))['properties']['time']
    for w in d['forecast']:
        for b in w['bins']:
            if b['probability'] > 1e-3:
                pr = rj_prob(P['a'], P['b'], P['p'], P['c'], P.get('aSigma') or 0, P['magMain'], b['magnitude'], (w['timeStart'] - tm) / 864e5, (w['timeEnd'] - tm) / 864e5)
                res.setdefault(m['name'].split('(')[-1].rstrip(')'), []).append(abs(pr - b['probability']) / b['probability'])
for k, v in res.items(): print(f'{k:22s} n={len(v):5d} median relative error {np.median(v):.3f}, 90th percentile {np.quantile(v, .9):.3f}')
