"""Registered scorer (protocol-kin2.json). Usage: score_kin2.py GSS.dta > results-kin2.json"""
import hashlib, json, sys
import numpy as np, pandas as pd
YEARS = (2000, 2010, 2018)
d = pd.read_stata(sys.argv[1], columns=["year", "age", "sex", "sibs", "childs", "born", "wtssps"], convert_categoricals=False)
d = d[d.year.isin(YEARS) & (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)]
text = json.load(open("text-kin2.json"))
res = {"protocolSha256": hashlib.sha256(open("protocol-kin2.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 = {q: {"engine": [], "text": [], "maternal": []} for q in ("sibs", "childs")}
for y in YEARS:
    eng = json.load(open(f"engine-kin2-{y}.json")); tx = text[str(y)]; res["years"][y] = {}
    for q, lim in (("sibs", 60), ("childs", 8)):
        dq = d[(d.year == y) & (d[q] <= lim)]
        surv, n = [], []
        for sx, b in cells:
            c = dq[(dq.sex == sx) & dq.age.between(*bands[b])]
            surv.append(round(float(np.average(c[q], weights=c.wtssps)), 3)); n.append(int(len(c)))
        e = eng[q][0] + eng[q][1]
        t = [round(float(np.mean([r[q][i] for r in tx])), 3) for i in range(6)]
        r = {"survey": surv, "n": n, "engine": e, "textModel": t,
             "engineMAE": round(float(np.mean(np.abs(np.array(e) - surv))), 3), "textModelMAE": round(float(np.mean(np.abs(np.array(t) - surv))), 3)}
        pool[q]["engine"] += list(np.abs(np.array(e) - surv)); pool[q]["text"] += list(np.abs(np.array(t) - surv))
        if q == "sibs":
            m = eng["sibsMaternal"][0] + eng["sibsMaternal"][1]; r["maternal"] = m
            r["maternalMAE"] = round(float(np.mean(np.abs(np.array(m) - surv))), 3); pool[q]["maternal"] += list(np.abs(np.array(m) - surv))
        res["years"][y][q] = r
res["pooled"] = {q: {k: round(float(np.mean(v)), 3) for k, v in pool[q].items() if v} for q in pool}
res["primary_met"] = bool(res["pooled"]["sibs"]["engine"] < res["pooled"]["sibs"]["text"])
print(json.dumps(res, indent=1))