"""Score the registered path-dependence test (protocol-squeeze.json) from the runs, one JSON line each.

    python3 experiment/path/score-squeeze.py <runs.jsonl> <protocol.json> <engine commit>
"""
import hashlib, json, statistics, sys

runs_path, protocol_path, commit = sys.argv[1:4]
runs = [json.loads(l) for l in open(runs_path) if l.strip()]
sizes = [1000, 10000, 100000, 1000000]

def summary(xs, sign):
    xs = [x for x in xs if x == x]
    return {"mean": round(statistics.fmean(xs), 3) if xs else None, "sd": round(statistics.stdev(xs), 3) if len(xs) > 1 else None,
            "predicted_sign": sum(1 for x in xs if x * sign > 0), "zero": sum(1 for x in xs if x == 0), "of": len(xs)}

by = {}
for n in sizes:
    rs = sorted((r for r in runs if r["N"] == n), key=lambda r: r["seed"])
    by[str(n)] = {
        "runs": len(rs),
        "gap_men": summary([r["gap_men"] for r in rs], -1), "gap_women": summary([r["gap_women"] for r in rs], 1),
        "direct_gap_men": summary([r["direct_gap_men"] for r in rs], -1), "direct_gap_women": summary([r["direct_gap_women"] for r in rs], 1),
        "ever_partnered_men_baseline_pct": round(statistics.fmean(100 * r["base"]["everMen"] / r["base"]["men"] for r in rs if r["base"]["men"]), 2),
        "direct_partnered_men_baseline_pct": round(statistics.fmean(100 * r["base"]["directMen"] / r["base"]["men"] for r in rs if r["base"]["men"]), 2),
        "cohort_men_women_baseline": [round(statistics.fmean(r["base"]["men"] for r in rs)), round(statistics.fmean(r["base"]["women"] for r in rs))],
        "cohort_men_women_shock": [round(statistics.fmean(r["shock"]["men"] for r in rs)), round(statistics.fmean(r["shock"]["women"] for r in rs))],
        "births_not_stored": sum(r["base"]["lost"] + r["shock"]["lost"] for r in rs),
        "seeds": [{"seed": r["seed"], "gap_men": round(r["gap_men"], 3), "gap_women": round(r["gap_women"], 3),
                   "direct_gap_men": round(r["direct_gap_men"], 3), "direct_gap_women": round(r["direct_gap_women"], 3)} for r in rs],
    }

big, small = by["1000000"], by["1000"]
P1 = big["runs"] == 20 and big["gap_men"]["predicted_sign"] == 20 and big["gap_women"]["predicted_sign"] == 20
P2 = small["runs"] == 20 and small["gap_men"]["predicted_sign"] < 18
stable = next((n for n in sizes if by[str(n)]["runs"] == 20 and by[str(n)]["gap_men"]["predicted_sign"] == 20), None)
out = {
    "protocolSha256": hashlib.sha256(open(protocol_path, "rb").read()).hexdigest(), "engineCommit": commit,
    "bySize": by,
    "decision": {"P1_met": P1, "P2_met": P2, "P3_smallest_size_all_20_seeds_men": stable},
    "note": "The direct answer's gap is zero in expectation, not exactly: the two arms hold different lives once the shock changes who is born, so the direct gaps reported are sampling noise around zero.",
}
print(json.dumps(out, indent=1))