"""Registered scorer for protocol-kin.json. Usage: score_kin.py <GSS .dta>"""
import sys, json, hashlib, numpy as np, pandas as pd
E = json.load(open("engine-2024.json")); T = json.load(open("text-2024.json"))
d = pd.read_stata(sys.argv[1], columns=["year", "age", "sex", "sibs", "childs", "born", "wtssps"], convert_categoricals=False)
d = d[(d.year == 2024) & (d.born == 1)]
bands = [(25, 34), (35, 44), (45, 59)]; cells = [(2, b) for b in range(3)] + [(1, b) for b in range(3)]
res = {"protocolSha256": hashlib.sha256(open("protocol-kin.json", "rb").read()).hexdigest(), "cells": ["women 25-34", "women 35-44", "women 45-59", "men 25-34", "men 35-44", "men 45-59"]}
for q, lim in (("childs", 8), ("sibs", 60)):
    truth, n = [], []
    for sx, b in cells:
        x = d[(d.sex == sx) & d.age.between(*bands[b]) & d[q].between(0, lim)]
        truth.append(float(np.average(x[q], weights=x.wtssps))); n.append(len(x))
    eng = [E[q][0][b] for b in range(3)] + [E[q][1][b] for b in range(3)]
    txt = list(np.mean([t[q] for t in T], 0))
    me, mt = float(np.mean(np.abs(np.array(eng) - truth))), float(np.mean(np.abs(np.array(txt) - truth)))
    res[q] = {"survey": [round(v, 3) for v in truth], "n": n, "engine": eng, "textModel": [round(v, 3) for v in txt], "engineMAE": round(me, 3), "textModelMAE": round(mt, 3)}
    print(q, "survey", [round(v, 2) for v in truth], "\n   engine", eng, "MAE", round(me, 3), "\n   text  ", [round(v, 2) for v in txt], "MAE", round(mt, 3))
res["primary"] = {"engineMAE": res["childs"]["engineMAE"], "textModelMAE": res["childs"]["textModelMAE"], "met": res["childs"]["engineMAE"] < res["childs"]["textModelMAE"]}
res["note"] = "The committed engine counts siblings as the mother's other children plus the father's children with other mothers (a development change), not the mother's children only as the protocol text says; this affects only the secondary siblings measure."
json.dump(res, open("results-kin.json", "w"), indent=1); print(json.dumps(res["primary"]))