"""Registered scorer (protocol-kin4.json). Usage: score_kin4.py GSS.dta > results-kin4.json"""
import hashlib, json, sys
import numpy as np, pandas as pd
YEARS = (2002, 2006, 2014)
d = pd.read_stata(sys.argv[1], columns=["year", "age", "sex", "sibs", "born", "wtssps"], convert_categoricals=False)
d = d[d.year.isin(YEARS) & (d.born == 1) & (d.sibs <= 60)]
bands = [(25, 34), (35, 44), (45, 59)]; cells = [(2, b) for b in range(3)] + [(1, b) for b in range(3)]
text = json.load(open("text-kin4.json"))
res = {"protocolSha256": hashlib.sha256(open("protocol-kin4.json", "rb").read()).hexdigest(), "cells": ["women 25-34", "women 35-44", "women 45-59", "men 25-34", "men 35-44", "men 45-59"], "years": {}}
pool = {"engine": [], "textModel": [], "kin2": [], "maternal": []}
for y in YEARS:
    eng = json.load(open(f"engine-kin4-{y}.json")); old = json.load(open(f"engine-kin2-{y}.json")); tx = text[str(y)]
    surv, n = [], []
    for sx, b in cells:
        c = d[(d.year == y) & (d.sex == sx) & d.age.between(*bands[b])]
        surv.append(round(float(np.average(c.sibs, weights=c.wtssps)), 3)); n.append(int(len(c)))
    e = eng["sibs"][0] + eng["sibs"][1]; k2 = old["sibs"][0] + old["sibs"][1]; mat = eng["maternal"][0] + eng["maternal"][1]
    t = [round(float(np.mean([r["sibs"][i] for r in tx])), 3) for i in range(6)]
    r = {"survey": surv, "n": n, "engine": e, "textModel": t, "kin2": k2, "maternal": mat}
    for k, v in (("engine", e), ("textModel", t), ("kin2", k2), ("maternal", mat)):
        err = np.abs(np.array(v) - surv); r[k + "MAE"] = round(float(err.mean()), 3); pool[k] += list(err)
    res["years"][y] = r
res["pooled"] = {k: round(float(np.mean(v)), 3) for k, v in pool.items()}
res["primary_met"] = bool(res["pooled"]["engine"] < res["pooled"]["textModel"] and res["pooled"]["engine"] < res["pooled"]["kin2"])
print(json.dumps(res, indent=1))