"""
Supporting numbers for the out-of-sample test of Davidson, U-Yen & Holloman (2015): window coverage
by polar phase, half-minicycle counts, the 40 uT amplitude test, each earthquake's distance to the
nearest window, the phase-stratified test, the polar-maximum-only test and the power calculations.

Needs, in the same folder as this script and spf_windows.py:
  wso_polar.html  - saved from http://wso.stanford.edu/Polar.html
  usgs_m75.csv    - saved from https://earthquake.usgs.gov/fdsnws/event/1/query?format=csv&starttime=1976-01-01&endtime=2026-10-06&minmagnitude=7.5&orderby=time-asc
Uses the same frozen rules ({'c13': 'ge9'}) as davidson_out_of_sample.py.
"""
import datetime as dt, numpy as np
from math import comb
from spf_windows import *

WSO = 'wso_polar.html'
quakes = load_quakes('usgs_m75.csv')
events = delay_events(quakes)
FROZEN = {'c13': 'ge9'}
def sf(k, n, p): return sum(comb(n, i) * p**i * (1-p)**(n-i) for i in range(k, n+1))


def analyse(data_end, day0, day1, label):
    dates, N, S, nint = load_wso(WSO, end=data_end)
    inst, ph, poles, min_ends = window_instances(dates, N, S, FROZEN)
    cov = build_coverage(dates, inst, events, day0, day1)
    days = [day0 + dt.timedelta(days=i) for i in range(len(cov))]
    ridx = np.clip(np.searchsorted(dates, days, side='right') - 1, 0, len(ph) - 1)
    dph = ph[ridx]
    print(f"\n### {label}: {day0}..{day1} ({len(cov)} days), WSO {dates[0]}..{dates[-1]} ({len(N)} points)")
    print(f"coverage {cov.mean()*100:.2f}%  | by phase: " + ", ".join(
        f"{nm} {cov[dph==k].mean()*100:.1f}% of {int((dph==k).sum())} d ({(dph==k).mean()*100:.0f}% of days)"
        for k, nm in ((MAX, 'max'), (MIN, 'min'), (REC, 'rec')) if (dph == k).any()))
    print(f"half-minicycles: N {len(poles['N'][1])}, S {len(poles['S'][1])}")
    for name in ('N', 'S'):
        X, hcs = poles[name]
        tot = ok = 0
        for q in range(1, len(hcs) - 1):
            h = hcs[q]
            if not (day0 <= dates[h['start']] <= day1):
                continue
            tot += 1
            ok += (abs(X[min(h['P'])] - X[min(hcs[q-1]['P'])]) >= 40 and abs(X[min(h['P'])] - X[min(hcs[q+1]['P'])]) >= 40)
        print(f"  {name}: half-cycles starting in period {tot}, passing the 40 uT amplitude test {ok} ({ok/max(tot,1)*100:.0f}%)")
    print(f"  max |N|, |S| in period: {np.abs(N[(dates >= day0) & (dates <= day1)]).max():.0f}, {np.abs(S[(dates >= day0) & (dates <= day1)]).max():.0f} uT")
    tgt = [q for q in quakes if q[1] >= 8.0 and day0 <= q[0] <= day1]
    idx = np.where(cov)[0]
    rows = []
    for d, m, p in tgt:
        k = (d - day0).days
        dist = int(np.min(np.abs(idx - k))) if len(idx) else None
        rows.append((d, m, ['MAX', 'MIN', 'REC'][dph[k]], bool(cov[k]), dist, p))
    return cov, dph, rows


cov_i, dph_i, rows_i = analyse(dt.date(2014, 1, 17), dt.date(1976, 7, 5), dt.date(2014, 1, 17), 'IN-SAMPLE')
print("  in-sample quakes 2004-2013:", sum(1 for r in rows_i if 2004 <= r[0].year <= 2013), "of", len(rows_i))
for r in rows_i:
    if r[0] in (dt.date(2004, 12, 26), dt.date(2005, 3, 28), dt.date(2007, 9, 12), dt.date(2010, 2, 27), dt.date(2011, 3, 11), dt.date(2012, 4, 11)):
        print("   fig4:", r[0], r[1], r[2], 'IN' if r[3] else 'out', f"{r[4]} d")
cov_o, dph_o, rows_o = analyse(None, dt.date(2014, 1, 18), dt.date(2026, 1, 9), 'OUT-OF-SAMPLE')
for r in rows_o:
    print(f"   {r[0]} M{r[1]} {r[2]} {'inside' if r[3] else 'outside'} nearest window {r[4]} d | {r[5][:40]}")

# phase-stratified chance probabilities (coverage of the phase each quake fell in)
phcov = {nm: cov_o[dph_o == k].mean() for k, nm in ((MAX, 'MAX'), (MIN, 'MIN'), (REC, 'REC'))}
probs = [phcov[r[2]] for r in rows_o]
dist = [1.0]
for pi in probs:
    new = [0.0] * (len(dist) + 1)
    for kk, v in enumerate(dist):
        new[kk] += v * (1 - pi); new[kk + 1] += v * pi
    dist = new
print("\nphase coverage out of sample:", {k: round(float(v) * 100, 1) for k, v in phcov.items()})
print("phase-stratified expected hits:", round(sum(probs), 2), " P(>=4):", round(sum(dist[4:]), 3))
nmax = sum(1 for r in rows_o if r[2] == 'MAX'); kmax = sum(1 for r in rows_o if r[2] == 'MAX' and r[3])
print(f"polar maximum only: {kmax} of {nmax}, P(>= {kmax} | {phcov['MAX']:.3f}) = {sf(kmax, nmax, phcov['MAX']):.3f}")
p = cov_o.mean()
crit = min(k for k in range(10) if sf(k, 9, p) <= 0.05)
print(f"hits needed for p<0.05: {crit};  power if 78.8% real: {sf(crit, 9, 0.788):.3f}")
for p1 in (0.6, 0.5):
    for n in range(1, 400):
        ks = [k for k in range(n + 1) if sf(k, n, p) <= 0.05]
        if ks and sf(min(ks), n, p1) >= 0.8:
            print(f"true rate {p1}: 80% power needs {n} quakes"); break