"""Registered scorer (protocol-floor.json). Scores the records model against the
sampling floor of each survey year: a prediction passes when its cell error is no worse than what a
model that knew the true population shares would score 95 times in 100. Two anchors: the year's own
headline share (subgroup split given the national level) and last survey's headline share (no
information from the target year)."""
import sys, numpy as np, pandas as pd
def bisect(f, lo, hi):
    for _ in range(80):
        mid = (lo + hi) / 2
        if f(lo) * f(mid) <= 0: hi = mid
        else: lo = mid
    return (lo + hi) / 2
import json
P = json.load(open("protocol-floor.json")); Q = {q: [1, 2, 3] for q in P["questions"]}; YEARS = [P["target_year"]]
cols = ["year", "age", "sex", "race", "degree", "partyid", "polviews", "region", "wtssps"] + list(Q)
d = pd.read_stata(sys.argv[1], columns=cols, convert_categoricals=False)
d = d[d["age"].between(18, 89) & d["sex"].isin([1, 2]) & d["race"].isin([1, 2, 3]) & d["degree"].isin(range(5)) & d["partyid"].isin(range(7)) & d["polviews"].isin(range(1, 8)) & d["region"].isin(range(1, 10)) & d["wtssps"].notna()]
def X(df):
    a = (df["age"].clip(18, 89).to_numpy(float) - 45) / 17; c = [np.ones(len(df)), a, a ** 2, (df["sex"] == 1).astype(float).to_numpy()]
    c += [(df["race"] == k).astype(float).to_numpy() for k in (2, 3)] + [(df["degree"] == k).astype(float).to_numpy() for k in range(1, 5)]
    c += [(df["partyid"] == k).astype(float).to_numpy() for k in range(1, 7)] + [(df["polviews"] == k).astype(float).to_numpy() for k in (1, 2, 3, 5, 6, 7)]
    c += [(df["region"] == k).astype(float).to_numpy() for k in range(2, 10)]
    p = (df["partyid"].to_numpy() <= 2) * 1.0 - (df["partyid"].to_numpy() >= 4)
    c += [p * a]
    return np.column_stack(c)
def fit(M, y, sw, l2=1e-3, iters=4000, lr=0.5):
    b = np.zeros(M.shape[1]); sw = sw / sw.mean()
    for _ in range(iters): p = 1 / (1 + np.exp(-(M @ b))); b -= lr * (M.T @ (sw * (p - y)) / len(y) + l2 * b)
    return b
pg = lambda p: np.where(p <= 2, 0, np.where(p == 3, 1, 2)); ag = lambda a: np.where(a < 45, 0, 1)
rng = np.random.default_rng(7)
print("question  year  floor95 | own-topline  pass | last-topline pass")
tally = [0, 0, 0]
for q, codes in Q.items():
    dq = d[d[q].isin(codes)].copy(); dq["y"] = (dq[q] == 1).astype(float)
    for T in YEARS:
        tr, te = dq[dq.year < T], dq[dq.year == T]
        half = 4.0  # recency: weight halves every 4 years
        b = fit(X(tr), tr.y.to_numpy(), 0.5 ** ((T - tr.year.to_numpy()) / half)); z = X(te) @ b
        w = te.wtssps.to_numpy(); y = te.y.to_numpy(); c = pg(te.partyid.to_numpy()) * 2 + ag(te.age.to_numpy())
        prev = tr[tr.year == tr.year.max()]
        out = []
        for top in (np.average(y, weights=w), np.average(prev.y, weights=prev.wtssps)):
            s = bisect(lambda s: np.average(1 / (1 + np.exp(-(z + s))), weights=w) - top, -10, 10); p = 1 / (1 + np.exp(-(z + s)))
            out.append(np.mean([abs(np.average(p[c == k], weights=w[c == k]) - np.average(y[c == k], weights=w[c == k])) for k in range(6)]) * 100)
        obs = np.array([np.average(y[c == k], weights=w[c == k]) for k in range(6)]); maes = []
        for _ in range(400):
            ix = np.concatenate([rng.choice(np.flatnonzero(c == k), (c == k).sum()) for k in range(6)])
            maes.append(np.mean([abs(np.average(y[ix][c[ix] == k], weights=w[ix][c[ix] == k]) - obs[k]) for k in range(6)]) * 100)
        f95 = np.quantile(maes, 0.95); tally[0] += 1; tally[1] += out[0] <= f95; tally[2] += out[1] <= f95
        print(f"{q:9s} {T:5d} {f95:7.2f} | {out[0]:11.2f} {'yes' if out[0] <= f95 else 'no':>5s} | {out[1]:12.2f} {'yes' if out[1] <= f95 else 'no':>4s}")
print(f"passed: own-topline {tally[1]}/{tally[0]}, last-topline {tally[2]}/{tally[0]}")
need = P["primary"]["questions_needed"]
print(json.dumps({"primary_met": bool(tally[1] >= need), "own_topline_passed": int(tally[1]), "last_topline_passed": int(tally[2]), "of": tally[0], "needed": need}))