"""The sealed rule for "more babies are born at the full moon".

Committed and pushed before any birth count was fetched (see README.md). The
function `analyse` is the whole rule; `branches.py` runs it on synthetic
series, `analyse.py` on the real ones.

Input: a DataFrame with columns `date` (datetime.date) and `births` (daily
count). Output: a dict with the estimate, its interval, the verdict, the two
calendar positive controls and the detection floor.

Method, in words:
1. Full-moon instants from `ephem`, converted to a calendar date at a fixed
   UTC-6 (US Central standard time; a birth is recorded on its local date).
2. The window is the full-moon date and the day after.
3. Days removed from the fit: US federal holidays as `holidays.US` lists them
   (observed dates included) and the day after each, 24 December - 2 January,
   14 February, 31 October, 29 February.
4. log(births) on fixed effects for year x month and year x weekday, by least
   squares on the kept days. Residuals for every day from that fit.
5. Estimate: mean residual on window days minus mean residual on other kept
   days, reported as a percentage (exp - 1).
6. Standard error by placebo windows: in each lunation, independently, a
   two-day window at a random offset at least two days clear of the true
   windows; the same difference computed against days in neither window. SE
   is the standard deviation of those placebo estimates. 95 % interval is
   estimate +/- 1.96 SE, on the log scale.
7. Verdict: survived if the lower bound is above 0; killed otherwise if the
   upper bound is below +0.5 %; inconclusive otherwise.
8. Positive controls, a gate: 14 February and 31 October against the other
   days within a week either side in the same year, SE by one random placebo
   day a year from that fortnight. If the instrument does not find
   Valentine's Day above zero AND Halloween below zero at 95 %, the study is
   void (the instrument cannot see a calendar effect known to exist).
9. Floor: the smallest true excess the rule would call survived four times in
   five, (1.96 + 0.84) SE.
"""

import datetime as dt

import ephem
import holidays
import numpy as np
import pandas as pd

TZ_HOURS = -6
WINDOW = (0, 1)
KILL_BAR = np.log(1.005)
PLACEBOS = 2000
Z = 1.959964
Z80 = 0.841621


def full_moon_dates(first, last):
    d = ephem.Date(dt.datetime(first.year, first.month, first.day) - dt.timedelta(days=45))
    out = []
    while True:
        d = ephem.next_full_moon(d)
        local = (d.datetime() + dt.timedelta(hours=TZ_HOURS)).date()
        if local > last + dt.timedelta(days=45):
            return out
        out.append(local)
        d = ephem.Date(d + 1)


def excluded_dates(first, last):
    years = range(first.year - 1, last.year + 2)
    ex = set()
    for h in holidays.US(years=years):
        ex.add(h)
        ex.add(h + dt.timedelta(days=1))
    for y in years:
        for day in range(24, 32):
            ex.add(dt.date(y, 12, day))
        ex.add(dt.date(y, 1, 1))
        ex.add(dt.date(y, 1, 2))
        ex.add(dt.date(y, 2, 14))
        ex.add(dt.date(y, 10, 31))
        if y % 4 == 0 and (y % 100 != 0 or y % 400 == 0):
            ex.add(dt.date(y, 2, 29))
    return ex


def calendar(dates):
    """Moon offset, lunation id and exclusion flag for each date. Depends only
    on the dates, so branches.py can reuse it across simulated series."""
    dates = list(dates)
    moons = full_moon_dates(dates[0], dates[-1])
    ex = excluded_dates(dates[0], dates[-1])
    mo = np.array([np.datetime64(m) for m in moons])
    dn = np.array([np.datetime64(d) for d in dates])
    idx = np.searchsorted(mo, dn, side="right") - 1
    offset = (dn - mo[idx]).astype(int)
    length = (mo[idx + 1] - mo[idx]).astype(int)
    keep = np.array([d not in ex for d in dates])
    return {"dates": dates, "offset": offset, "lunation": idx, "length": length, "keep": keep}


def residuals(dates, births):
    s = pd.Series(dates)
    year = s.map(lambda d: d.year).to_numpy()
    month = s.map(lambda d: d.month).to_numpy()
    wday = s.map(lambda d: d.weekday()).to_numpy()
    ym = pd.factorize(year * 100 + month)[0]
    yw = pd.factorize(year * 10 + wday)[0]
    n = len(dates)
    X = np.zeros((n, ym.max() + 1 + yw.max() + 1))
    X[np.arange(n), ym] = 1.0
    X[np.arange(n), ym.max() + 1 + yw] = 1.0
    return X


def fit(cal, births):
    y = np.log(np.asarray(births, dtype=float))
    X = residuals(cal["dates"], births)
    k = cal["keep"]
    beta, *_ = np.linalg.lstsq(X[k], y[k], rcond=None)
    return y - X @ beta


def _diff(r, inside, outside):
    return r[inside].mean() - r[outside].mean()


def moon_estimate(cal, r, rng, placebos=PLACEBOS):
    k = cal["keep"]
    win = np.isin(cal["offset"], WINDOW)
    est = _diff(r, win & k, ~win & k)
    lun = cal["lunation"]
    ids, first = np.unique(lun, return_index=True)
    lengths = cal["length"][first]
    pos = np.searchsorted(ids, lun)
    ests = np.empty(placebos)
    for i in range(placebos):
        # offsets 3 .. L-3 keep a two-day window at least two days from the
        # true windows at 0-1 and L-L+1
        o = rng.integers(3, lengths - 2)
        start = o[pos]
        p = (cal["offset"] == start) | (cal["offset"] == start + 1)
        ests[i] = _diff(r, p & k, ~p & ~win & k)
    se = ests.std(ddof=1)
    return est, se


def day_control(cal, r, month, day, rng, placebos=PLACEBOS):
    dates = cal["dates"]
    years = sorted({d.year for d in dates})
    date_index = {d: i for i, d in enumerate(dates)}
    targets, neighbours = [], []
    for y in years:
        t = dt.date(y, month, day)
        if t not in date_index:
            continue
        nb = [date_index[t + dt.timedelta(days=j)] for j in range(-7, 8)
              if j != 0 and (t + dt.timedelta(days=j)) in date_index]
        nb = [i for i in nb if cal["keep"][i]]
        if len(nb) < 7:
            continue
        targets.append(date_index[t])
        neighbours.append(nb)
    est = np.mean([r[t] - r[nb].mean() for t, nb in zip(targets, neighbours)])
    ests = np.empty(placebos)
    for i in range(placebos):
        vals = []
        for nb in neighbours:
            j = rng.integers(len(nb))
            others = [x for m, x in enumerate(nb) if m != j]
            vals.append(r[nb[j]] - r[others].mean())
        ests[i] = np.mean(vals)
    return est, ests.std(ddof=1), len(targets)


def verdict(est, se):
    lo, hi = est - Z * se, est + Z * se
    if lo > 0:
        return "survived"
    if hi < KILL_BAR:
        return "killed"
    return "inconclusive"


def pct(x):
    return 100.0 * (np.exp(x) - 1.0)


def analyse(df, seed=20260911, placebos=PLACEBOS, controls=True):
    df = df.sort_values("date").reset_index(drop=True)
    cal = calendar(df["date"])
    r = fit(cal, df["births"])
    rng = np.random.default_rng(seed)
    est, se = moon_estimate(cal, r, rng, placebos)
    out = {
        "days": int(len(df)),
        "kept": int(cal["keep"].sum()),
        "window_days": int((np.isin(cal["offset"], WINDOW) & cal["keep"]).sum()),
        "residual_sd": float(r[cal["keep"]].std(ddof=1)),
        "est": est,
        "se": se,
        "est_pct": pct(est),
        "lo_pct": pct(est - Z * se),
        "hi_pct": pct(est + Z * se),
        "floor_pct": pct((Z + Z80) * se),
        "verdict": verdict(est, se),
    }
    if controls:
        v, vse, vn = day_control(cal, r, 2, 14, rng, placebos)
        h, hse, hn = day_control(cal, r, 10, 31, rng, placebos)
        out["valentine"] = (pct(v), pct(v - Z * vse), pct(v + Z * vse), vn)
        out["halloween"] = (pct(h), pct(h - Z * hse), pct(h + Z * hse), hn)
        out["controls_pass"] = bool(v - Z * vse > 0 and h + Z * hse < 0)
        if not out["controls_pass"]:
            out["verdict"] = "void: the positive controls failed"
        curve = pd.Series(r[cal["keep"]]).groupby(cal["offset"][cal["keep"]]).mean()
        out["curve_pct"] = {int(k): float(pct(v2)) for k, v2 in curve.items()}
    out["cal"] = cal
    return out