"""Build goal -> tactic training pairs from the CoqGym corpus (Yang & Deng 2019).

Nothing from dalil al-kalima or dalil al-imkan is used: the only source is the
third-party CoqGym corpus.  ssreflect-style projects are dropped because our
proofs use plain Ltac.  Output: data/sft_{train,valid}.jsonl  with fields
prompt / completion / project / family.
"""
import gzip, json, random, re, sys
from pathlib import Path
HERE = Path(__file__).resolve().parent
SRC = HERE.parents[1] / "app/static/research/training/files"
OUT = HERE / "data"; OUT.mkdir(exist_ok=True)
DROP_PROJECTS = {"math-comp", "smc", "fourcolor", "odd-order"}
SSR = re.compile(r"^\s*(move|case:|elim:|apply/|have\b|rewrite\s*/|by\s*\[|done|congr|suff)|=>|//|/=")
MAXG, MAXT = 1500, 160

def prompt(goal):
    return f"Goal:\n{goal.strip()}\nTactic:"

def load(split, n=None, seed=0):
    rows, seen = [], set()
    for line in gzip.open(SRC / f"coqgym_{split}.jsonl.gz", "rt"):
        r = json.loads(line)
        g, t = r["goal_before"], r["tactic"].strip()
        if r["project"] in DROP_PROJECTS or not g or g.rstrip().endswith("=" * 28) or len(g) > MAXG or len(t) > MAXT or SSR.search(t):
            continue
        t = re.sub(r"\s+", " ", t)
        k = (g, t)
        if k in seen: continue
        seen.add(k)
        rows.append({"prompt": prompt(g), "completion": " " + t, "project": r["project"], "family": r["family"]})
    random.Random(seed).shuffle(rows)
    return rows[:n] if n else rows

if __name__ == "__main__":
    tr = load("train"); va = load("valid", 2000)
    with open(OUT / "sft_train.jsonl", "w") as f:
        for r in tr: f.write(json.dumps(r) + "\n")
    with open(OUT / "sft_valid.jsonl", "w") as f:
        for r in va: f.write(json.dumps(r) + "\n")
    print(len(tr), len(va))