"""Registered scorer for protocol-voice-split.json. Usage: score_voice_split.py <GSS .dta>"""
import json, sys, glob, hashlib, numpy as np, pandas as pd
P = json.load(open("protocol-voice-split.json")); FRESH = list(P["questions"])
DEV = ["gunlaw", "cappun", "grass", "natenvir", "natheal", "natrace", "natcrime", "nateduc", "natarms", "natsoc", "natspac"]
load = lambda pat: [json.loads(l) for f in sorted(glob.glob(pat)) for l in open(f)]
Vd = load("dev-cells-flash.jsonl"); Vf = load("fresh-cells-gemini-3.8-flash.jsonl")
cods = lambda q: [1, 2] if q in ("gunlaw", "cappun", "grass") else [1, 2, 3]
d = pd.read_stata(sys.argv[1], columns=["year", "age", "partyid", "wtssps"] + DEV + FRESH, convert_categoricals=False)
d = d[d.year.isin([2022, 2024]) & d.age.between(18, 89) & d.partyid.isin(range(7)) & d.wtssps.notna()]
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)
def row(q, T, V):
    x = d[(d.year == T) & d[q].isin(cods(q))]; y = (x[q] == 1).to_numpy(float); w = x.wtssps.to_numpy(); c = pg(x.partyid.to_numpy()) * 2 + ag(x.age.to_numpy())
    obs = np.array([np.average(y[c == k], weights=w[c == k]) for k in range(6)]) * 100; top = np.average(y, weights=w) * 100
    cw = np.array([w[c == k].sum() for k in range(6)]); cw = cw / cw.sum(); 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]) * 100 - obs[k]) for k in range(6)]))
    vs = [v["answer"]["cells"] for v in V if v["q"] == q and v["year"] == T and v["answer"] and len(v["answer"].get("cells", [])) == 6]
    return dict(q=q, T=T, obs=obs, top=top, cw=cw, vc=np.mean(vs, 0), asks=len(vs), f95=float(np.quantile(maes, 0.95)))
lg = lambda p: np.log(p / (100 - p)); ex = lambda z: 100 / (1 + np.exp(-z))
def bis(f):
    lo, hi = -10, 10
    for _ in range(60):
        m = (lo + hi) / 2
        if f(lo) * f(m) <= 0: hi = m
        else: lo = m
    return (lo + hi) / 2
def lev(r, lam):
    z = lg(np.clip(r["vc"], 1, 99)); z = z.mean() + lam * (z - z.mean()); s = bis(lambda s: (r["cw"] * ex(z + s)).sum() - r["top"]); return ex(z + s)
mae = lambda p, r: float(np.mean(np.abs(p - r["obs"])))
dev = [row(q, T, Vd) for q in DEV for T in (2022, 2024)]
grid = np.round(np.arange(0.30, 1.3001, 0.05), 2); lam = float(grid[np.argmin([np.mean([mae(lev(o, g), o) for o in dev]) for g in grid])])
out = []
for q in FRESH:
    r = row(q, 2024, Vf); v, b = mae(lev(r, lam), r), mae(np.full(6, r["top"]), r)
    out.append({"question": q, "asks": r["asks"], "floor95": round(r["f95"], 2), "voice": round(v, 2), "baseline": round(b, 2), "insideFloor": v <= r["f95"],
                "observed": [round(x, 1) for x in r["obs"]], "voiceCells": [round(x, 1) for x in lev(r, lam)]})
    print(f"{q:9s} floor {r['f95']:.2f} voice {v:.2f} baseline {b:.2f} inside {v <= r['f95']}")
mv, mb = np.mean([o["voice"] for o in out]), np.mean([o["baseline"] for o in out]); inside = sum(o["insideFloor"] for o in out)
res = {"protocolSha256": hashlib.sha256(open("protocol-voice-split.json", "rb").read()).hexdigest(), "lambda": lam, "questions": out,
       "primary": {"meanVoice": round(float(mv), 2), "meanBaseline": round(float(mb), 2), "insideFloor": int(inside), "of": 5,
                   "met": bool(mv < mb and inside >= 3)}}
json.dump(res, open("results-voice-split.json", "w"), indent=1); print(json.dumps(res["primary"]), "lambda", lam)