"""Checks run AFTER the data, labelled that way wherever they are quoted.

    python mate.py > mate.txt

1. The two controls on raw counts, with no fitted model: each 31 October (and
   14 February) against the mean of the same weekday one week before and one
   week after. A weekday-free ratio, so it cannot be an artefact of the fit.
2. Injection on the real series: multiply window-day births by 1.003 and by
   1.005 and rerun the sealed rule, to show on the real data what the rule does
   with an excess of known size (the branch table did this on synthetic noise).
3. The two files where they overlap, 2000-2003: how closely the SSA and NCHS
   daily counts agree.
"""

import datetime as dt

import numpy as np
import pandas as pd

import rule
from analyse import load

ssa = load("US_births_2000-2014_SSA.csv")
nchs_all = load("US_births_1994-2003_CDC_NCHS.csv")

print("== 1. controls on raw counts: the day / mean(same weekday -7 and +7 days)")
by_date = dict(zip(ssa["date"], ssa["births"]))
for label, (m, d) in (("31 Oct", (10, 31)), ("14 Feb", (2, 14))):
    ratios = []
    for y in range(2000, 2015):
        t = dt.date(y, m, d)
        a, b = t - dt.timedelta(days=7), t + dt.timedelta(days=7)
        ratios.append(by_date[t] / ((by_date[a] + by_date[b]) / 2) - 1)
    r = 100 * np.array(ratios)
    print(f"{label}: mean {r.mean():+.2f} %, median {np.median(r):+.2f} %, "
          f"range {r.min():+.2f} to {r.max():+.2f} %, years below zero {int((r < 0).sum())} of 15")
    print("   by year: " + " ".join(f"{y}:{v:+.1f}" for y, v in zip(range(2000, 2015), r)))
    days = [dt.date(y, m, d).strftime("%a") for y in range(2000, 2015)]
    weekend = np.array([dt.date(y, m, d).weekday() >= 5 for y in range(2000, 2015)])
    print("   weekday: " + " ".join(f"{y}:{w}" for y, w in zip(range(2000, 2015), days)))
    print(f"   weekend years mean {r[weekend].mean():+.2f} % (n {int(weekend.sum())}), "
          f"weekday years mean {r[~weekend].mean():+.2f} % (n {int((~weekend).sum())})")

print()
print("== 2. a known excess injected on window days of the real SSA series")
cal = rule.calendar(sorted(ssa["date"]))
win = np.isin(cal["offset"], rule.WINDOW)
base = ssa.sort_values("date").reset_index(drop=True)
for excess in (0.003, 0.005):
    df = base.copy()
    df["births"] = df["births"] * np.where(win, 1 + excess, 1.0)
    out = rule.analyse(df, controls=False)
    print(f"+{100*excess:.1f} % injected: estimate {out['est_pct']:+.3f} % "
          f"({out['lo_pct']:+.3f} to {out['hi_pct']:+.3f}), verdict {out['verdict']}")

print()
print("== 3. SSA against NCHS on the days both cover, 2000-2003")
both = ssa.merge(nchs_all, on="date", suffixes=("_ssa", "_nchs"))
ratio = both["births_ssa"] / both["births_nchs"]
print(f"days {len(both)}; SSA / NCHS mean {ratio.mean():.4f}, "
      f"sd {ratio.std():.4f}, min {ratio.min():.4f}, max {ratio.max():.4f}; "
      f"correlation {np.corrcoef(both['births_ssa'], both['births_nchs'])[0, 1]:.4f}")