"""Join forecasts.csv with ComCat counts. Writes outcomes.csv: one row per (version, window, M) whose window ended
before CUTOFF, with observed count N, the stated P(N>=1), median, 95% range, and a randomized PIT from the fractiles."""
import sys; sys.path.insert(0, '/work/aftershock-forecasts')
from common import *
CUTOFF = datetime(2026, 9, 26, tzinfo=timezone.utc).timestamp() * 1000   # catalog fetched 2026-10-03; allow 1 week lag
FR = np.round(np.arange(0.01, 0.99001, 0.005), 3)
V = {v['vid']: v for v in csv.DictReader(open('versions.csv'))}
rng = np.random.default_rng(0)
GEN = {}   # event -> generic RJ parameters from its earliest 'Generic' version
for v in sorted(V.values(), key=lambda v: int(v['creation_ms'])):
    if 'Generic' in v['model'] and v['eid'] not in GEN and v['a']:
        GEN[v['eid']] = tuple(float(v[k] or 0) for k in ('a', 'b', 'p', 'c', 'aSigma'))
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):
    """Reasenberg-Jones P(N>=1) in [t0,t1] days after the mainshock, integrating over a ~ N(a, asig)."""
    I = ((t1 + c) ** (1 - p) - (t0 + c) ** (1 - p)) / (1 - p) if abs(p - 1) > 1e-9 else np.log((t1 + c) / (t0 + c))
    lam = 10 ** (a + asig * Z + b * (Mm - M)) * I
    return float(1 - (WZ * np.exp(-lam)).sum()), float(10 ** (a + b * (Mm - M)) * I)
import collections
MC = collections.defaultdict(collections.Counter)   # event -> most common stated Mcat across its versions
for v in V.values():
    if v['Mcat']: MC[v['eid']][float(v['Mcat'])] += 1
MCAT = {e: c.most_common(1)[0][0] for e, c in MC.items()}
BG = {}
def bg_rate(eid, rlat, rlon, rkm, M, tm):
    """Events/day with mag>=M-0.05 in the circle over the 10 years before the mainshock (None if not fetched)."""
    p = f'raw/bg/{eid}.csv'
    if not os.path.exists(p): return None
    if eid not in BG:
        R = [r for r in csv.DictReader(open(p)) if r['mag']]
        BG[eid] = (np.array([float(r['latitude']) for r in R]), np.array([float(r['longitude']) for r in R]), np.array([float(r['mag']) for r in R]))
    la, lo, mg = BG[eid]
    return float(((mg >= M - 0.05) & (hav(rlat, rlon, la, lo) <= rkm)).sum()) / 3652.5
out = []
for f in csv.DictReader(open('forecasts.csv')):
    v = V[f['vid']]
    if not v['rkm'] or int(f['t1_ms']) > CUTOFF or catalog(v['eid']) is None: continue
    M = float(f['M']); N = count(v['eid'], float(v['rlat']), float(v['rlon']), float(v['rkm']), int(f['t0_ms']), int(f['t1_ms']), M)
    p = float(f['prob'])
    fv = np.array(f['fractiles'].split(), float) if f['fractiles'] else None
    def F(n):  # predictive CDF P(count <= n)
        if n < 0: return 0.0
        if n == 0: return 1 - p
        if fv is None or len(fv) != len(FR): return np.nan
        k = np.nonzero(fv <= n)[0]
        return max(1 - p, FR[k[-1]] if len(k) else 0.0)
    lo, hi = F(N - 1), F(N)
    pit = lo + (hi - lo) * rng.random() if np.isfinite(hi) else np.nan
    tm = int(v['time_ms']); d0 = (int(f['t0_ms']) - tm) / 864e5; d1 = (int(f['t1_ms']) - tm) / 864e5
    g = GEN.get(v['eid']); pg = rj_prob(*g, float(v['magMain'] or v['mag']), M, d0, d1)[0] if g else np.nan
    br = bg_rate(v['eid'], float(v['rlat']), float(v['rlon']), float(v['rkm']), M, tm)
    pbg = 1 - np.exp(-br * (d1 - d0)) if br is not None else np.nan
    out.append(dict(d0=round(d0, 4), d1=round(d1, 4), p_generic=round(pg, 6), p_bg=round(pbg, 6), vid=f['vid'], eid=v['eid'], mag=v['mag'], net=v['net'], Mcat=v['Mcat'], Mcat_ev=MCAT.get(v['eid'], ''), model=v['model'], rkm=v['rkm'],
                    creation_ms=v['creation_ms'], time_ms=v['time_ms'], window=f['window'], t0_ms=f['t0_ms'], t1_ms=f['t1_ms'],
                    M=M, prob=p, median=f['median'], lo95=f['lo95'], hi95=f['hi95'], N=N, pit=round(pit, 4)))
with open('outcomes.csv', 'w', newline='') as fh:
    w = csv.DictWriter(fh, fieldnames=list(out[0])); w.writeheader(); w.writerows(out)
print(len(out), 'rows,', len({o['eid'] for o in out}), 'events')