#!/usr/bin/env python3
"""
Check a Living Worlds trace without the engine.

Every chance in the run is decided by draw(seed, life, month, event), a public 32-bit hash, compared
with the published monthly tables. This script rebuilds the world from population.csv and replays,
in the same order the engine uses, every death, birth, baby's sex and separation, and checks that the
events in events.csv are exactly those and no others. Partnerships are matched by the engine; for
those it checks that both partners were alive, unpartnered, old enough, in the same place, and that
the woman's own draw fell under her chance of seeking a partner that month.

  python3 check_trace.py <trace directory>

Standard library only. Exits 0 when every check passes.
"""
import csv, hashlib, json, os, sys

M32 = 0xFFFFFFFF
def imul(a, b): return (a * b) & M32
def draw(seed, life, month, event):
    h = (seed ^ imul(life, 0x9E3779B1) ^ imul(month, 0x85EBCA77) ^ imul(event, 0xC2B2AE3D)) & M32
    h = imul(h ^ (h >> 16), 0x7FEB352D)
    h = imul(h ^ (h >> 15), 0x846CA68B)
    h ^= h >> 16
    return h / 4294967296.0
DEATH, BIRTH, SEX, PARTNER, SEPARATE = 1, 2, 3, 4, 5
AGES, BIRTH_AGES, PARITIES, MAX_AGE = 121, 56, 8, 120

def main(d):
    man = json.load(open(os.path.join(d, "manifest.json")))
    for f, want in man["files"].items():
        got = hashlib.sha256(open(os.path.join(d, f), "rb").read()).hexdigest()
        if got != want: sys.exit(f"FAIL {f}: sha256 {got} does not match the manifest")
    seed, months = man["seed"], man["months"]
    hz = json.load(open(os.path.join(d, "hazards.json")))
    born, male, place, died, parity, partner, mother = [], [], [], [], [], [], []
    for r in csv.DictReader(open(os.path.join(d, "population.csv"))):
        born.append(int(r["born"])); male.append(r["sex"] == "M"); place.append(int(r["place"]))
        died.append(None); parity.append(0); partner.append(None); mother.append(None)
    logged = {}
    for r in csv.DictReader(open(os.path.join(d, "events.csv"))):
        logged.setdefault(int(r["month"]), []).append((r["event"], int(r["life"]), r["other"]))
    totals = {"death": 0, "birth": 0, "sex": 0, "separation": 0, "partnership": 0}
    for m in range(months):
        want = {"death": set(), "separation": set(), "birth": []}
        for i in range(len(born)):
            if died[i] is not None or born[i] > m: continue
            age = min(MAX_AGE, (m - born[i]) // 12)
            if draw(seed, i, m, DEATH) < hz["death"][(1 if male[i] else 0) * AGES + age]:
                died[i] = m; want["death"].add(i)
                p = partner[i]
                if p is not None: partner[p] = None; partner[i] = None
                continue
            if not male[i] and partner[i] is not None and draw(seed, i, m, SEPARATE) < hz["separate"][age]:
                p = partner[i]; partner[p] = None; partner[i] = None; want["separation"].add((i, p))
            if not male[i] and 12 <= age < BIRTH_AGES and draw(seed, i, m, BIRTH) < hz["birth"][age * PARITIES + min(PARITIES - 1, parity[i])]:
                want["birth"].append((i, "M" if draw(seed, i, m, SEX) < hz["boys"] else "F"))
        got = logged.get(m, [])
        gd = {l for e, l, o in got if e == "death"}
        gs = {(l, int(o)) for e, l, o in got if e == "separation"}
        gb = [(int(o.split(":")[0]), o.split(":")[1], l) for e, l, o in got if e == "birth"]
        if gd != want["death"]: sys.exit(f"FAIL month {m}: deaths differ ({len(gd ^ want['death'])} lives)")
        if gs != want["separation"]: sys.exit(f"FAIL month {m}: separations differ")
        if [(a, b) for a, b, _ in gb] != want["birth"]: sys.exit(f"FAIL month {m}: births or a baby's sex differ")
        for mo, sx, child in gb:
            if child != len(born): sys.exit(f"FAIL month {m}: child id {child}, expected {len(born)}")
            born.append(m + 1); male.append(sx == "M"); place.append(place[mo]); died.append(None)
            parity.append(0); partner.append(None); mother.append(mo); parity[mo] += 1
        for e, w, o in got:
            if e != "partnership": continue
            mn = int(o); aw = (m - born[w]) // 12; am = (m - born[mn]) // 12
            ok = (not male[w] and male[mn] and died[w] is None and died[mn] is None and partner[w] is None and partner[mn] is None
                  and aw >= 16 and am >= 16 and place[w] == place[mn] and draw(seed, w, m, PARTNER) < hz["partner"][min(MAX_AGE, aw)])
            if not ok: sys.exit(f"FAIL month {m}: partnership {w}-{mn} breaks a rule")
            partner[w] = mn; partner[mn] = w
        totals["death"] += len(gd); totals["separation"] += len(gs); totals["birth"] += len(gb); totals["sex"] += len(gb)
        totals["partnership"] += sum(1 for e, _, _ in got if e == "partnership")
    print(json.dumps({"result": "PASS", "lives_at_start": man["lives"], "months": months, "seed": seed,
                      "recomputed_exactly": {k: totals[k] for k in ("death", "birth", "sex", "separation")},
                      "partnerships_checked_against_rules": totals["partnership"]}))

if __name__ == "__main__": main(sys.argv[1] if len(sys.argv) > 1 else ".")