import json, hashlib, collections

R4  = "/home/infra-ops/erp-tune/recipe-r4/survivors-r4.jsonl"
GOV = "/home/infra-ops/erp-tune/datasets/derived/govreport/clean-v1/govreport-clean-v1.jsonl"
QMS = "/home/infra-ops/erp-tune/datasets/derived/qmsum/clean-v1/qmsum-clean-v1.jsonl"
OUT = "/home/infra-ops/erp-tune/recipe-r5/survivors-r5.jsonl"

# Start from run-4's survivors: they already carry the EXACT held dialogue
# selection (c2-logs 301 / cwm 183 / bluemoon 126), fireball whole (5,396) and
# kvasir's exact 1,613-sample prefix cut. Reusing them is what guarantees kvasir
# is byte-identical to run 4 -- NOT re-cut (run-5 recipe held clause).
r4 = [json.loads(l) for l in open(R4)]
rows = [r for r in r4 if r["dataset_id"] != "airoboros-3.2"]

c = collections.Counter(r["dataset_id"] for r in rows)
held = {"c2-logs-32k-llama3": 301, "creative-writing-multiturn": 183,
        "bluemoon": 126, "fireball": 5396, "kvasir": 1613}
assert set(c) == set(held), "unexpected dataset_id set after airoboros removal: %s" % dict(c)
for k, v in held.items():
    assert c[k] == v, "HELD ROOT CHANGED: %s %d != %d" % (k, c[k], v)
print("  held (dialogue+fireball+kvasir) unchanged:", dict(c))

def append_root(path, dsid, expect):
    n = 0
    seen = set()
    for l in open(path):
        d = json.loads(l)
        assert d["id"] not in seen, "dup id in %s: %s" % (dsid, d["id"])
        seen.add(d["id"])
        rows.append({"dataset_id": dsid, "id": d["id"]})
        n += 1
    assert n == expect, "%s: %d != %d" % (dsid, n, expect)
    print("  %s appended: %d" % (dsid, n))

append_root(GOV, "govreport", 496)
append_root(QMS, "qmsum", 97)

expected_total = 301 + 183 + 126 + 5396 + 1613 + 496 + 97  # 8212
assert len(rows) == expected_total, "%d != %d" % (len(rows), expected_total)

with open(OUT, "w") as f:
    for r in rows:
        f.write(json.dumps(r, separators=(",", ":"), sort_keys=True) + "\n")
sha = hashlib.sha256(open(OUT, "rb").read()).hexdigest()
print("  wrote %s survivors -> %s" % (format(len(rows), ","), OUT))
print("  final composition:", dict(collections.Counter(r["dataset_id"] for r in rows)))
print("  sha256 %s" % sha)