"""Registered scorer for protocol-prior.json. Usage: score_prior.py <GSS .dta>"""
import sys, json, hashlib, numpy as np, pandas as pd
P = json.load(open("protocol-prior.json")); Q = list(P["questions"])
V = [json.loads(l) for l in open("prior-cells-gemini-3.8-flash.jsonl")]
d = pd.read_stata(sys.argv[1], columns=["year", "age", "partyid", "wtssps"] + Q, convert_categoricals=False)
d = d[(d.year == 2024) & d.age.between(18, 89) & d.partyid.isin(range(7)) & d.wtssps.notna()]
codes = {"nataid": [1, 2, 3], "conjudge": [1, 2, 3], "confed": [1, 2, 3], "immfate": [1, 2, 3], "polhitok": [1, 2]}
lg = lambda p: np.log(p / (1 - p)); ex = lambda z: 1 / (1 + np.exp(-z))
def bis(f):
    lo, hi = -10, 10
    for _ in range(50):
        m = (lo + hi) / 2
        if f(lo) * f(m) <= 0: hi = m
        else: lo = m
    return (lo + hi) / 2
rng = np.random.default_rng(20261005); K = 50; arms = ["model+50", "real50", "real75", "real100"]; allerr = {a: [] for a in arms}; out = []
for q in Q:
    x = d[d[q].isin(codes[q])]; y = (x[q] == 1).to_numpy(float); w = x.wtssps.to_numpy()
    c = np.where(x.partyid <= 2, 0, np.where(x.partyid == 3, 1, 2)) * 2 + np.where(x.age < 45, 0, 1)
    obs = np.array([np.average(y[c == k], weights=w[c == k]) for k in range(6)]); cw = np.array([w[c == k].sum() for k in range(6)]); cw /= cw.sum()
    fl = []
    for _ in range(300):
        ix = np.concatenate([rng.choice(np.flatnonzero(c == k), (c == k).sum()) for k in range(6)])
        fl.append(np.mean([abs(np.average(y[ix][c[ix] == k], weights=w[ix][c[ix] == k]) - obs[k]) for k in range(6)]))
    f95 = float(np.quantile(fl, 0.95))
    vs = [v["answer"]["cells"] for v in V if v["q"] == q and v["answer"] and len(v["answer"].get("cells", [])) == 6]
    z0 = lg(np.clip(np.mean(vs, 0) / 100, 0.01, 0.99)); z0 = z0.mean() + 0.85 * (z0 - z0.mean())
    samp = lambda n: np.array([y[rng.choice(np.flatnonzero(c == k), n, p=w[c == k] / w[c == k].sum())].mean() for k in range(6)])
    E = {a: [] for a in arms}
    for _ in range(400):
        ph = samp(50); top = (cw * ph).sum(); s = bis(lambda s: (cw * ex(z0 + s)).sum() - top); prior = ex(z0 + s)
        E["model+50"].append(np.mean(np.abs((K * prior + 50 * ph) / (K + 50) - obs))); E["real50"].append(np.mean(np.abs(ph - obs)))
        E["real75"].append(np.mean(np.abs(samp(75) - obs))); E["real100"].append(np.mean(np.abs(samp(100) - obs)))
    for a in arms: allerr[a] += E[a]
    out.append({"question": q, "asks": len(vs), "floor95": round(f95 * 100, 2), **{a: {"meanError": round(float(np.mean(E[a])) * 100, 2), "insideFloor": round(float(np.mean(np.array(E[a]) <= f95)), 3)} for a in arms}})
    print(q, {a: round(float(np.mean(E[a])) * 100, 2) for a in arms}, "floor", round(f95 * 100, 2))
pooled = {a: round(float(np.mean(allerr[a])) * 100, 3) for a in arms}
res = {"protocolSha256": hashlib.sha256(open("protocol-prior.json", "rb").read()).hexdigest(), "questions": out, "pooledMeanError": pooled,
       "primary": {"modelPlus50": pooled["model+50"], "real75": pooled["real75"], "met": pooled["model+50"] < pooled["real75"]},
       "secondary": {"beatsReal100": pooled["model+50"] < pooled["real100"]}}
json.dump(res, open("results-prior.json", "w"), indent=1); print(json.dumps(res["primary"]), json.dumps(res["secondary"]), pooled)