"""Run the HTPS search on a problem set and verify every proof with coqc.

Usage: python run_eval.py --set synth:data/synth_test.jsonl|heldout --proposer fixed|model:PATH|hybrid:PATH
       --tag NAME [--workers 6] [--expansions 48] [--seconds 120] [--limit N]
Writes results/<tag>.jsonl (one line per problem, including the search log)
and prints a summary.
"""
import argparse, json, re, sys, threading, time
from pathlib import Path
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE)); sys.path.insert(0, str(HERE.parent / "pilot"))
from rocq_env import Coqtop, Goal, verify, with_hammer
import htps
import gen_problems

RES = HERE / "results"; RES.mkdir(exist_ok=True)


def problem_set(spec):
    if spec == "heldout":
        import pilot   # only the vocabulary preamble and the task STATEMENTS -- never any proof
        defs = re.findall(r"^Definition (\w+)", pilot.PREAMBLE, re.M)
        cfg = dict(prelude=pilot.PREAMBLE, ambient={"Thing", "Gives", "Through", "Exists", "Ess", "PartOf", "Universe"},
                   defs=defs, consts={"Thing": ["Universe"]}, section_end="End Pilot.",
                   preds={"Gives": 2, "Through": 2, "Exists": 1, "Ess": 1, "PartOf": 2, "DependsOn": 2,
                          "ThroughAnother": 1, "OfItself": 1, "Necessary": 1, "Possible": 1, "InUniverse": 1})
        probs = [{"id": k, "stmt": v} for k, v in pilot.TASKS.items()]
        return cfg, probs
    kind, path = spec.split(":", 1)
    cfg = dict(prelude=gen_problems.PRELUDE, ambient=gen_problems.AMBIENT, defs=list(gen_problems.DEFS),
               consts={"T": ["a", "b"]}, section_end="End Synth.",
               preds={"P": 1, "Q": 1, "R": 1, "E": 2, "F": 2, "A": 0, "B": 0,
                      **{d: ar for d, (ar, _) in gen_problems.DEFS.items() if ar > 0}})
    probs = [json.loads(l) for l in open(HERE / path)]
    return cfg, probs


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--set", required=True); ap.add_argument("--proposer", required=True)
    ap.add_argument("--tag", required=True); ap.add_argument("--workers", type=int, default=6)
    ap.add_argument("--expansions", type=int, default=48); ap.add_argument("--seconds", type=float, default=120)
    ap.add_argument("--k", type=int, default=8); ap.add_argument("--case-splits", action="store_true"); ap.add_argument("--critic", action="store_true"); ap.add_argument("--max-asserts", type=int, default=None); ap.add_argument("--sample", type=int, default=0); ap.add_argument("--hammer", action="store_true"); ap.add_argument("--assert-splits", action="store_true"); ap.add_argument("--eager-leaves", action="store_true"); ap.add_argument("--ids", default=""); ap.add_argument("--decompose-only", action="store_true"); ap.add_argument("--temperature", type=float, default=0.8); ap.add_argument("--limit", type=int, default=0)
    a = ap.parse_args()
    cfg, probs = problem_set(a.set)
    if a.limit: probs = probs[:a.limit]
    if a.ids: probs = [p for p in probs if p["id"] in a.ids.split(",")]
    # start the Rocq processes before loading any model (forking during model load is unsafe)
    coqs = [Coqtop(cfg["prelude"], cfg["ambient"], hammer=a.hammer) for _ in range(a.workers)]
    if a.proposer == "fixed":
        prop = htps.FixedProposer(cfg["defs"], cfg["consts"], cfg.get("preds") if a.case_splits else None)
    elif a.proposer.startswith("hybrid:"):
        prop = htps.HybridProposer(htps.ModelProposer(a.proposer.split(":", 1)[1], beams=a.k, sample=a.sample, temperature=a.temperature),
                                   htps.FixedProposer(cfg["defs"], cfg["consts"], cfg.get("preds") if a.case_splits else None))
    else:
        prop = htps.ModelProposer(a.proposer.split(":", 1)[1], beams=a.k, sample=a.sample, temperature=a.temperature)
    out = open(RES / f"{a.tag}.jsonl", "w")
    lock = threading.Lock()
    todo = list(probs)
    stats = {"n": 0, "proved": 0, "verified": 0}
    t0 = time.time()

    def worker(coq):
        while True:
            with lock:
                if not todo: return
                p = todo.pop(0)
            log = []
            r = htps.search(coq, Goal([], p["stmt"]), prop, cfg["defs"], k=a.k,
                            max_expansions=a.expansions, max_seconds=a.seconds, log=log, critic=a.critic, max_asserts=a.max_asserts,
                            hammer=a.hammer, decompose_only=a.decompose_only,
                            assert_splits=a.assert_splits, eager_leaves=a.eager_leaves)
            r["verified"] = False
            if r["proved"]:
                ok, vout, src = verify(with_hammer(cfg["prelude"]) if a.hammer else cfg["prelude"], "problem", p["stmt"],
                                       r["script"], cfg["section_end"], hammer=a.hammer)
                r["verified"], r["verify_output"] = ok, vout[-1500:]
            r.pop("tree", None)
            rec = {"id": p["id"], "stmt": p["stmt"], "proposer": prop.name, **r, "log": log}
            with lock:
                stats["n"] += 1; stats["proved"] += r["proved"]; stats["verified"] += r["verified"]
                out.write(json.dumps(rec) + "\n"); out.flush()
                print(f"[{stats['n']}/{len(probs)}] {p['id']}: {'VERIFIED' if r['verified'] else ('proved-but-REJECTED' if r['proved'] else 'no')} "
                      f"exp={r['expansions']} t={r['seconds']}s  running {stats['verified']}/{stats['n']}", flush=True)

    ths = [threading.Thread(target=worker, args=(c,)) for c in coqs]
    for t in ths: t.start()
    for t in ths: t.join()
    summ = {"tag": a.tag, "set": a.set, "proposer": prop.name, "problems": len(probs), **stats,
            "expansions": a.expansions, "seconds_limit": a.seconds, "k": a.k, "wall_seconds": round(time.time() - t0)}
    print(json.dumps(summ))
    with open(RES / "summary.jsonl", "a") as f:
        f.write(json.dumps(summ) + "\n")


if __name__ == "__main__":
    main()