"""Freeze the kernel learn/probe split (run once, before any model run on the probe set).

Validity rules (preregistered; reference properties only, from the strict-fp32 A100 survey, never model results):
  ok on A100-40GB; finite outputs; not all zero; cache_margin >= 10 (an input-independent answer would break the
  1e-4 tolerance by >= 10x, so caching/constant outputs cannot pass); reference re-run difference < 1e-5
  (deterministic enough for the tolerance); peak memory <= 10 GiB; inputs + outputs <= 2 GiB; eager time in
  [0.02, 100] ms.
Split: seeded (SEED) and disjoint by problem id. Probe = 6 valid level-2 + 4 valid level-3 problems (headroom);
learn = every other valid problem (levels 1-3). Writes splits.json with the probe hash (harness.api.probe_hash).

usage: python3 -m domains.kernel.make_splits domains/kernel/evidence/survey_a100_strict.jsonl
"""
from __future__ import annotations

import json
import os
import random
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.normpath(os.path.join(HERE, "..", "..")))

from harness.api import TaskRef, probe_hash  # noqa: E402

SEED = 20261008
N_PROBE = {"L2": 6, "L3": 4}
GIB = 2 ** 30


def valid(r: dict) -> tuple[bool, str]:
    if not r.get("ok"):
        return False, "reference failed: " + str(r.get("error", ""))[:80]
    if not r.get("finite"):
        return False, "non-finite output"
    if r.get("all_zero"):
        return False, "all-zero output"
    if float(r.get("cache_margin", 0)) < 10:
        return False, f"cache_margin {r.get('cache_margin'):.3g} < 10 (input-independent answer would pass)"
    if float(r.get("rerun_maxdiff", 1)) >= 1e-5:
        return False, f"nondeterministic reference (rerun diff {r.get('rerun_maxdiff'):.3g})"
    if int(r.get("peak_mem_bytes", 0)) > 10 * GIB:
        return False, "peak memory > 10 GiB"
    if int(r.get("input_bytes", 0)) + int(r.get("output_bytes", 0)) > 2 * GIB:
        return False, "inputs + outputs > 2 GiB"
    t = float(r.get("eager_ms", 0))
    if not 0.02 <= t <= 100:
        return False, f"eager time {t:.3g} ms out of range"
    return True, ""


def main():
    rows = [json.loads(ln) for ln in open(sys.argv[1])]
    verdict = {r["pid"]: valid(r) for r in rows}
    ok = sorted(p for p, (v, _w) in verdict.items() if v)
    rnd = random.Random(SEED)
    probe = []
    for lvl, n in N_PROBE.items():
        pool = [p for p in ok if p.startswith(lvl)]
        probe += sorted(rnd.sample(pool, n))
    learn = [p for p in ok if p not in probe]
    refs = [TaskRef("kernel", p, "", "probe") for p in probe]
    out = {"version": 1, "seed": SEED, "rules": __doc__.split("Validity rules")[1].split("Split:")[0].strip(),
           "survey": os.path.basename(sys.argv[1]), "probe": probe, "learn": learn,
           "probe_hash": probe_hash(refs),
           "excluded": {p: w for p, (v, w) in sorted(verdict.items()) if not v},
           "counts": {"valid": len(ok), "learn": len(learn), "probe": len(probe),
                      "learn_by_level": {lv: sum(p.startswith(lv) for p in learn) for lv in ("L1", "L2", "L3")}}}
    json.dump(out, open(os.path.join(HERE, "splits.json"), "w"), indent=1)
    print(json.dumps({k: out[k] for k in ("probe", "probe_hash", "counts")}, indent=1))


if __name__ == "__main__":
    main()