"""Score protocol-gunlaw.json. The only code that opens sealed-2024.pkl.
    python score.py <datadir> <protocol.json> <voices-2022.jsonl> <voices-2024.jsonl> <arm name>"""
import hashlib, json, sys, numpy as np, pandas as pd
D, P, V22, V24, ARM = sys.argv[1:6]
rng = np.random.default_rng(20261005)
train = pd.read_pickle(f"{D}/train.pkl"); test = pd.read_pickle(f"{D}/sealed-2024.pkl").merge(pd.read_pickle(f"{D}/test-covariates.pkl"), on="id")
def X(df):
    a = (df["age"].clip(18, 89).to_numpy(float) - 45) / 17; cols = [np.ones(len(df)), a, a ** 2, (df["sex"] == 1).astype(float).to_numpy()]
    cols += [(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)]
    cols += [(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)]
    cols += [(df["region"] == k).astype(float).to_numpy() for k in range(2, 10)]
    return np.column_stack(cols)
def fit(M, y, l2=1e-3, iters=4000, lr=0.5):
    w = np.zeros(M.shape[1])
    for _ in range(iters): p = 1 / (1 + np.exp(-(M @ w))); w -= lr * (M.T @ (p - y) / len(y) + l2 * w)
    return w
ytr = (train["answer"] == 1).astype(float).to_numpy(); W = fit(X(train), ytr)
Mte = X(test); z_reg = Mte @ W; P_reg = 1 / (1 + np.exp(-z_reg))
load = lambda p: {r["id"]: r["answer"] for r in map(json.loads, open(p)) if r.get("answer") in (1, 2)}
v22, v24 = load(V22), load(V24)
test["voice"] = test["id"].map(v24); ok = test["voice"].notna().to_numpy()
anc = train[train["id"].isin(v22)].copy(); va = np.where(anc["id"].map(v22) == 1, 1.0, -1.0); za = X(anc) @ W; ya = (anc["answer"] == 1).astype(float).to_numpy()
bs = np.linspace(0, 4, 81)
nll = lambda b: -np.mean(ya * np.log(1 / (1 + np.exp(-(za + b * va))) + 1e-12) + (1 - ya) * np.log(1 - 1 / (1 + np.exp(-(za + b * va))) + 1e-12))
b = float(bs[np.argmin([nll(x) for x in bs])])
vt = np.zeros(len(test)); vt[ok] = np.where(test.loc[ok, "voice"] == 1, 1.0, -1.0)
P_anc = 1 / (1 + np.exp(-(z_reg + b * vt))); P_anc[~ok] = P_reg[~ok]
pg = lambda p: np.where(p <= 2, "Democrat", np.where(p == 3, "Independent", "Republican"))
ag = lambda a: np.where(a < 45, "18-44", "45+")
test["pg"] = pg(test["partyid"].to_numpy()); test["ag"] = ag(test["age"].to_numpy())
t22 = train[train["year"] == 2022].copy(); t22["pg"] = pg(t22["partyid"].to_numpy()); t22["ag"] = ag(t22["age"].to_numpy())
cells = [(g, a) for g in ("Democrat", "Independent", "Republican") for a in ("18-44", "45+")]
cf = {c: float(np.average(t22.loc[(t22.pg == c[0]) & (t22.ag == c[1]), "answer"] == 1, weights=t22.loc[(t22.pg == c[0]) & (t22.ag == c[1]), "wtssps"])) for c in cells}
METHODS = ["carry_forward", "regression", "voice", "anchored"]
def shares(idx):
    t = test.iloc[idx]; w = t["wtssps"].to_numpy(); fav = (t["answer"].to_numpy() == 1); vo = t["voice"].to_numpy(); okk = ok[idx]; out = {}
    for c in cells:
        m = ((t.pg == c[0]) & (t.ag == c[1])).to_numpy(); vm = m & okk
        out[c] = {"measured": np.average(fav[m], weights=w[m]), "carry_forward": cf[c], "regression": np.average(P_reg[idx][m], weights=w[m]),
                  "anchored": np.average(P_anc[idx][m], weights=w[m]), "voice": np.average(vo[vm] == 1, weights=w[vm]) if vm.any() else np.nan}
    return out
def summary(sh):
    mae = {k: float(np.mean([abs(sh[c][k] - sh[c]["measured"]) for c in cells])) * 100 for k in METHODS}
    gaps = {g: {k: float(sh[(g, "18-44")][k] - sh[(g, "45+")][k]) * 100 for k in METHODS + ["measured"]} for g in ("Democrat", "Independent", "Republican")}
    return mae, gaps
sh = shares(np.arange(len(test))); mae, gaps = summary(sh)
boot = {k: [] for k in METHODS}; dvr = []; dar = []; gapb = {g: [] for g in gaps}
for _ in range(2000):
    idx = rng.integers(0, len(test), len(test)); m2, g2 = summary(shares(idx))
    for k in METHODS: boot[k].append(m2[k])
    dvr.append(m2["voice"] - m2["regression"]); dar.append(m2["anchored"] - m2["regression"])
    for g in gaps: gapb[g].append(g2[g]["measured"])
ci = lambda a: [float(np.nanpercentile(a, 2.5)), float(np.nanpercentile(a, 97.5))]
gap_ok = {k: all(ci(gapb[g])[0] <= gaps[g][k] <= ci(gapb[g])[1] for g in gaps) for k in METHODS}
y = (test["answer"].to_numpy() == 1).astype(float)
ll = lambda p, m: float(-np.mean(y[m] * np.log(p[m] + 1e-12) + (1 - y[m]) * np.log(1 - p[m] + 1e-12)))
allm = np.ones(len(test), bool); pv = np.where(vt > 0, 0.96, 0.04)
res = {"arm": ARM, "protocolSha256": hashlib.sha256(open(P, "rb").read()).hexdigest(), "N": int(len(test)), "voiceDidNotRun": int((~ok).sum()), "anchorB": b,
       "cells": {f"{c[0]} / {c[1]}": {k: round(float(v) * 100, 1) for k, v in sh[c].items()} for c in cells},
       "sixCellMAE": {k: {"points": round(mae[k], 2), "ci95": [round(x, 2) for x in ci(boot[k])]} for k in METHODS},
       "voiceMinusRegression": {"points": round(mae["voice"] - mae["regression"], 2), "ci95": [round(x, 2) for x in ci(dvr)]},
       "anchoredMinusRegression": {"points": round(mae["anchored"] - mae["regression"], 2), "ci95": [round(x, 2) for x in ci(dar)]},
       "ageGapByParty": {g: {k: round(v, 1) for k, v in gaps[g].items()} | {"measured_ci95": [round(x, 1) for x in ci(gapb[g])]} for g in gaps},
       "individual": {"regression": {"log_loss": ll(P_reg, allm), "accuracy": float(np.mean((P_reg > .5) == (y == 1)))},
                      "anchored": {"log_loss": ll(P_anc, allm), "accuracy": float(np.mean((P_anc > .5) == (y == 1)))},
                      "voice": {"log_loss_on_voiced": ll(pv, ok), "accuracy_on_voiced": float(np.mean((vt[ok] > 0) == (y[ok] == 1)))}},
       "overall2024": round(float(np.average(test["answer"] == 1, weights=test["wtssps"])) * 100, 1),
       "decision": {k: {"mae_le_3": mae[k] <= 3.0, "gaps_inside_ci": gap_ok[k], "matches_measured_split": bool(mae[k] <= 3.0 and gap_ok[k])} for k in METHODS}
                   | {"voice_adds_over_regression": bool(ci(dvr)[1] < 0), "anchored_adds_over_regression": bool(ci(dar)[1] < 0)}}
json.dump(res, open(f"{D}/results-gunlaw-{ARM}.json", "w"), indent=1); print(json.dumps({k: res[k] for k in ("arm", "N", "voiceDidNotRun", "anchorB", "sixCellMAE", "voiceMinusRegression", "anchoredMinusRegression", "decision")}, indent=1))