# Build SFT set: correct samples on train problems greedy got wrong (rejection sampling, STaR-style),
# max 2 distinct correct traces per problem, plus a replay mix of correct greedy traces.
import json, random, sys
sys.path.insert(0, "."); from grade import load
g = load("out/train_greedy.jsonl"); s = load("out/train_samples.jsonl")
wrong = {r["item"] for r in g if not r["ok"]}
keep = {}
for r in s:
    if r["item"] in wrong and r["ok"] and (r["predicted_n"] or 0) < 2048:
        L = keep.setdefault(r["item"], [])
        if len(L) < 2 and r["text"] not in [x["text"] for x in L]: L.append(r)
fixed = [x for L in keep.values() for x in L]
right = [r for r in g if r["ok"] and (r["predicted_n"] or 0) < 2048]
random.Random(20).shuffle(right)
replay = right[:len(fixed)]
sft = fixed + replay; random.Random(1).shuffle(sft)
with open("data/sft.jsonl", "w", encoding="utf-8") as f:
    for r in sft: f.write(json.dumps({"problem": r["problem"], "completion": r["text"], "src": "fixed" if r in fixed else "replay"}) + "\n")
st = {"train_items": len(g), "greedy_correct": len(g) - len(wrong), "greedy_wrong": len(wrong),
      "wrong_items_with_a_correct_sample": len(keep), "fixed_traces": len(fixed), "replay_traces": len(replay), "sft_total": len(sft)}
json.dump(st, open("logs/sft_stats.json", "w"), indent=1); print(st)