"""Expert iteration: turn search results into policy training data.

Every goal the search proved (the root or any subgoal) gives one pair
(goal -> the tactic that proved it), each step already checked by Rocq.
Steps that are just the automation closers are skipped.
HOLD-OUT: results from the held-out set (dalil al-kalima / dalil al-imkan
tasks) are refused -- they are never training data -- and so are the
synthetic test problems.

Usage: python collect.py out.jsonl results/a.jsonl [results/b.jsonl ...]
"""
import json, re, sys
from pathlib import Path
from htps import CLOSERS

def is_closer(t):
    # the closers are tried at every goal anyway; the policy should learn the
    # decomposition steps, not re-learn automation
    return t in CLOSERS or t.startswith("unfold") and t.endswith("; firstorder")

HERE = Path(__file__).resolve().parent
TEST_IDS = {json.loads(l)["id"] for f in ("synth_test_hard.jsonl", "synth_test_easy.jsonl") for l in open(HERE / "data" / f)}
out, files = sys.argv[1], sys.argv[2:]
pairs, seen = [], set()
for f in files:
    for line in open(f):
        r = json.loads(line)
        if not re.fullmatch(r"s\d+_\d+", r["id"]):
            raise SystemExit(f"refusing {f}: problem {r['id']} is not a synthetic problem (hold-out rule)")
        if r["id"] in TEST_IDS:
            raise SystemExit(f"refusing {f}: problem {r['id']} is in the synthetic test set")
        for goal, tac in r.get("proved_steps", []):
            if is_closer(tac): continue
            k = (goal, tac)
            if k in seen: continue
            seen.add(k)
            pairs.append({"prompt": f"Goal:\n{goal}\nTactic:", "completion": " " + tac + ".", "project": "synth-ei", "family": None})
with open(out, "w") as fo:
    for p in pairs: fo.write(json.dumps(p) + "\n")
print(len(pairs), "pairs from", len(files), "files")