"""Score the branches before the seal: P(survived), P(killed), P(inconclusive)
for the rule in rule.py, under assumed truths, on synthetic series.

No birth data is used. The calendar is the real one (2000-01-01 .. 2014-12-31,
the SSA file's span, with the real full moons and the real exclusions), and the
log birth count on each day is

    year x month effect + year x weekday effect + noise + true excess on window days

The two fixed effects are absorbed exactly by the rule's fit, so only the noise
and the excess matter; the noise is Gaussian with the stated SD, iid, or AR(1)
with rho 0.3 where marked. This is also the positive control for the moon
estimate: an excess of known size is injected and the table says how often the
rule sees it.

    python branches.py > branches.txt
"""

import datetime as dt

import numpy as np

import rule

SIMS = 400
PLACEBOS = 400
rng = np.random.default_rng(20260911)

dates = [dt.date(2000, 1, 1) + dt.timedelta(days=i)
         for i in range((dt.date(2014, 12, 31) - dt.date(2000, 1, 1)).days + 1)]
cal = rule.calendar(dates)
win = np.isin(cal["offset"], rule.WINDOW)
X = rule.residuals(dates, None)
k = cal["keep"]
Xk = X[k]
pinv = np.linalg.pinv(Xk)
n = len(dates)


def noise(sd, rho):
    e = rng.normal(0.0, sd, n)
    if rho:
        out = np.empty(n)
        out[0] = e[0]
        for i in range(1, n):
            out[i] = rho * out[i - 1] + e[i] * np.sqrt(1 - rho * rho)
        return out
    return e


print(f"rule: 2000-2014 calendar, window = full-moon date and the day after, "
      f"{SIMS} synthetic series per row, {PLACEBOS} placebos each")
print(f"days {n}, kept {int(k.sum())}, window days kept {int((win & k).sum())}")
print(f"{'true excess':>11} {'noise sd':>9} {'rho':>4} {'P(survived)':>12} "
      f"{'P(killed)':>10} {'P(inconcl.)':>12} {'mean SE %':>10}")
for excess in (0.0, 0.002, 0.005, 0.02):
    for sd, rho in ((0.015, 0.0), (0.03, 0.0), (0.03, 0.3)):
        tally = {"survived": 0, "killed": 0, "inconclusive": 0}
        ses = []
        for _ in range(SIMS):
            y = noise(sd, rho) + np.where(win, np.log1p(excess), 0.0)
            r = y - X @ (pinv @ y[k])
            est, se = rule.moon_estimate(cal, r, rng, PLACEBOS)
            tally[rule.verdict(est, se)] += 1
            ses.append(se)
        print(f"{excess:>+11.1%} {sd:>9.3f} {rho:>4.1f} {tally['survived']/SIMS:>12.4f} "
              f"{tally['killed']/SIMS:>10.4f} {tally['inconclusive']/SIMS:>12.4f} "
              f"{100*np.mean(ses):>10.3f}")