"""Critic training data from search logs (HTPS-style value targets).

Positive: every goal the search proved (root or subgoal).
Negative: every goal the search expanded but did not prove.
The model learns   "Goal:\n<goal>\nProvable:"  ->  " yes" / " no".
Negatives are subsampled to match the positives.
HOLD-OUT: only synthetic TRAINING problems are used (same guards as collect.py).

Usage: python collect_critic.py out.jsonl results/a.jsonl [...]
"""
import json, random, re, sys
from pathlib import Path

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)}


def critic_prompt(goal):
    return f"Goal:\n{goal}\nProvable:"


def main():
    out, files = sys.argv[1], sys.argv[2:]
    pos, neg = set(), 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")
            proved = {g for g, _ in r.get("proved_steps", [])}
            expanded = {e["goal"] for e in r.get("log", [])}
            pos |= proved
            neg |= expanded - proved
    neg -= pos
    neg = sorted(neg); random.Random(0).shuffle(neg)
    neg = neg[:len(pos)]
    rows = [{"prompt": critic_prompt(g), "completion": " yes", "project": "critic", "family": None} for g in pos]
    rows += [{"prompt": critic_prompt(g), "completion": " no", "project": "critic", "family": None} for g in neg]
    random.Random(1).shuffle(rows)
    with open(out, "w") as fo:
        for r in rows: fo.write(json.dumps(r) + "\n")
    print(len(pos), "positive,", len(neg), "negative")


if __name__ == "__main__":
    main()