"""Neural-symbolic pilot: HTPS-style subgoal decomposition, a pretrained language
model as the proposer, Rocq as the checker.

The model (Claude, through isolated stateless Claude Code calls: no tools, no
files, no project memory, run from an empty folder) sees ONLY the vocabulary,
the approved postulates, and the statement to prove -- never the repository's
proofs of dalil al-kalima or dalil al-imkan, which are the held-out test.

For each goal:
  1. symbolic leaf attempt: plain Rocq automation, no model (this is also the
     symbolic-only baseline);
  2. the model proposes direct proofs; Rocq checks each;
  3. the model proposes decompositions into subgoals plus a combining proof;
     Rocq checks the combination with the subgoals as hypotheses, then each
     subgoal is proved recursively (AND); alternatives are tried in order (OR).
The final proof is assembled into one file with no Admitted/admit/Axiom, and
Print Assumptions must show at most excluded middle.

Every model call and every Rocq check is logged to runs/<task>.jsonl.
Usage:  python3 pilot.py [task ...]      (no argument: all tasks)
"""
import json, os, re, subprocess, sys, tempfile, time
from pathlib import Path

HERE = Path(__file__).resolve().parent
RUNS = HERE / "runs"; RUNS.mkdir(exist_ok=True)
MODEL = "sonnet"
MAX_DEPTH = 3
MAX_CALLS = 30
FORBIDDEN = re.compile(r"\b(Admitted|admit|Axiom|Parameter|Hypothesis|Abort|Variable|Section|End)\b")

PREAMBLE = r"""From Stdlib Require Import Classical Relations.

Section Pilot.
(* Vocabulary.  Nothing is assumed about any of these. *)
Variable Thing : Type.
Variable Gives : Thing -> Thing -> Prop.     (* Gives y x : y gives existence to x *)
Variable Through : Thing -> Thing -> Prop.   (* Through x y : x has its existence through y *)
Variable Exists : Thing -> Prop.             (* x exists *)
Variable Ess : Thing -> Prop.                (* x's existence is from its own nature *)
Variable PartOf : Thing -> Thing -> Prop.    (* p is a part of u *)
Variable Universe : Thing.                   (* the universe *)

(* Definitions. *)
Definition DependsOn (x y : Thing) : Prop := Gives y x.            (* x depends on y *)
Definition ThroughAnother (x : Thing) : Prop := exists y, Gives y x.
Definition OfItself (x : Thing) : Prop := Exists x /\ ~ ThroughAnother x.
Definition Necessary (x : Thing) : Prop := Exists x /\ Ess x.
Definition Possible (x : Thing) : Prop := Exists x /\ ~ Ess x.
Definition InUniverse (x : Thing) : Prop := x = Universe \/ PartOf x Universe.

(* Approved postulates, written as definitions: a theorem uses one only when
   its statement lists it as a condition. *)
Definition MP001 : Prop := forall x y, Gives y x -> ~ Through y x.   (* a giver does not have its existence through what it gives existence to *)
Definition MP002 : Prop := forall x y, Gives y x -> Through x y.     (* the receiver has its existence through the giver *)
Definition MP003 : Prop := forall x y z, Through x y -> Through y z -> Through x z.  (* existence-through is inherited *)
Definition MP008 : Prop := forall x, Exists x -> ~ ThroughAnother x -> Ess x.       (* existence comes from one's nature or from a giver *)
Definition OP001 : Prop := forall x y, Gives y x -> OfItself y.      (* only what exists of itself gives existence *)
Definition OP006 : Prop := exists x, InUniverse x /\ Possible x.     (* the universe is not wholly necessary *)
Definition SP001 : Prop := exists x, Exists x.                       (* something exists *)
"""

UNFOLD = ("unfold MP001, MP002, MP003, MP008, OP001, OP006, SP001, OfItself, ThroughAnother, "
          "Necessary, Possible, InUniverse, DependsOn in *")
AUTOMATION = [
    "intros; firstorder.",
    f"{UNFOLD}; intros; firstorder.",
    f"{UNFOLD}; intros; intuition eauto.",
    f"{UNFOLD}; intros; tauto.",
]

TASKS = {
 "pr002_no_self_creation": "MP001 -> MP002 -> forall x, ~ Gives x x",
 "pr004_no_two_step_giving": "OP001 -> forall x y z, Gives y x -> ~ Gives z y",
 "pr003_no_circle": "MP001 -> MP002 -> MP003 -> forall x, ~ clos_trans Thing DependsOn x x",
 "step3_possible_needs_cause": "MP008 -> forall x, Possible x -> ThroughAnother x",
 "pr011_something_caused": "MP008 -> OP006 -> exists x, InUniverse x /\\ Possible x /\\ ThroughAnother x",
 "kalima_first_cause": "OP001 -> ThroughAnother Universe -> exists b, Gives b Universe /\\ OfItself b /\\ (forall y, ~ Gives y b)",
 "imkan_b_necessary_exists": "MP008 -> OP001 -> SP001 -> exists n, Necessary n",
 "pr001_relational_exhaustive": ("forall x, (exists y, (x = y \\/ clos_trans Thing DependsOn x y) /\\ DependsOn y y) \\/ "
    "(exists y, (x = y \\/ clos_trans Thing DependsOn x y) /\\ clos_trans Thing DependsOn y y /\\ ~ DependsOn y y) \\/ "
    "(exists S : Thing -> Prop, S x /\\ forall a, S a -> exists b, S b /\\ DependsOn a b) \\/ "
    "((forall y, ~ DependsOn x y) \\/ exists y, clos_trans Thing DependsOn x y /\\ forall z, ~ DependsOn y z)"),
}

SYSTEM = ("You are an expert in the Rocq Prover (version 9, formerly Coq), using its standard library. "
          "You write short, correct Rocq tactic proofs. You reply with JSON only, no prose, no markdown fences.")


class Run:
    def __init__(self, task):
        self.task, self.calls, self.log = task, 0, open(RUNS / f"{task}.jsonl", "w")
        self.lemmas = []          # proved lemmas: (name, statement, proof)
        self.counter = 0

    def record(self, **kw):
        kw["time"] = time.strftime("%Y-%m-%dT%H:%M:%S"); self.log.write(json.dumps(kw) + "\n"); self.log.flush()

    # ---------- Rocq
    def check(self, extra_hyps, lemmas, name, stmt, proof):
        body = PREAMBLE
        for hname, hstmt in extra_hyps:
            body += f"Hypothesis {hname} : {hstmt}.\n"
        for lname, lstmt, lproof in lemmas:
            body += f"\nLemma {lname} : {lstmt}.\nProof.\n{lproof}\nQed.\n"
        body += f"\nLemma {name} : {stmt}.\nProof.\n{proof}\nQed.\n\nEnd Pilot.\n"
        with tempfile.TemporaryDirectory() as d:
            f = Path(d) / "Check.v"; f.write_text(body)
            r = subprocess.run(["coqc", "-q", str(f)], capture_output=True, text=True, timeout=60)
        ok = r.returncode == 0
        return ok, (r.stdout + r.stderr)[-1500:]

    # ---------- model
    def ask(self, prompt):
        if self.calls >= MAX_CALLS:
            return None
        self.calls += 1
        with tempfile.TemporaryDirectory() as d:     # empty folder: no project, no memory
            t0 = time.time()
            r = subprocess.run(["claude", "-p", "--model", MODEL, "--tools", "", "--strict-mcp-config",
                                "--setting-sources", "", "--no-session-persistence", "--system-prompt", SYSTEM,
                                "--output-format", "json", prompt], cwd=d, capture_output=True, text=True, timeout=300)
        try:
            meta = json.loads(r.stdout); text = meta.get("result", "")
            model = list((meta.get("modelUsage") or {}).keys())
        except Exception:
            text, model = r.stdout, []
        self.record(kind="model_call", n=self.calls, model=model, seconds=round(time.time() - t0, 1), prompt=prompt, response=text)
        m = re.search(r"\{.*\}", text, re.S)
        try:
            return json.loads(m.group(0)) if m else None
        except Exception:
            return None

    def context(self):
        avail = "\n".join(f"Lemma {n} : {s}." for n, s, _ in self.lemmas) or "(none yet)"
        return (f"Rocq context (a Section; you cannot change it):\n{PREAMBLE}\n"
                f"Lemmas already proved in this section, which you may use by name:\n{avail}\n")

    # ---------- search
    def prove(self, stmt, depth, path):
        # 1. symbolic leaf attempt (no model)
        for tac in AUTOMATION:
            name = self.fresh()
            ok, out = self.check([], self.lemmas, name, stmt, tac)
            self.record(kind="check", stage="automation", path=path, stmt=stmt, proof=tac, ok=ok)
            if ok:
                self.lemmas.append((name, stmt, tac)); return name
        # 2. model: direct proofs
        resp = self.ask(self.context() + f"\nGoal to prove:\n{stmt}\n\nPropose up to 3 different complete tactic proofs "
                        "(the text between Proof. and Qed.). You may unfold definitions. Reply as JSON: "
                        '{"proofs": ["...", "...", "..."]}')
        for tac in (resp or {}).get("proofs", [])[:3]:
            if not isinstance(tac, str) or FORBIDDEN.search(tac):
                self.record(kind="rejected", reason="forbidden command or bad format", proof=tac); continue
            name = self.fresh()
            ok, out = self.check([], self.lemmas, name, stmt, tac)
            self.record(kind="check", stage="model_direct", path=path, stmt=stmt, proof=tac, ok=ok, rocq=out)
            if ok:
                self.lemmas.append((name, stmt, tac)); return name
        if depth >= MAX_DEPTH:
            return None
        # 3. model: decompositions into subgoals (AND), alternatives (OR)
        resp = self.ask(self.context() + f"\nGoal to prove:\n{stmt}\n\nIt was too hard to prove directly. Decompose it into "
                        "2 to 4 simpler subgoals (Rocq propositions over the section's vocabulary; each a closed statement), "
                        "and give a tactic proof of the goal that may use the subgoals as hypotheses named sub1, sub2, ... "
                        "Propose up to 2 alternative decompositions. Reply as JSON: "
                        '{"decompositions": [{"subgoals": ["...", "..."], "combine": "..."}]}')
        for dec in (resp or {}).get("decompositions", [])[:2]:
            subs, comb = dec.get("subgoals", []), dec.get("combine", "")
            if not subs or not isinstance(comb, str) or FORBIDDEN.search(comb) or any(FORBIDDEN.search(s) for s in subs):
                self.record(kind="rejected", reason="forbidden command or bad format", decomposition=dec); continue
            hyps = [(f"sub{i+1}", s) for i, s in enumerate(subs)]
            name = self.fresh()
            ok, out = self.check(hyps, self.lemmas, name, stmt, comb)
            self.record(kind="check", stage="model_combine", path=path, stmt=stmt, subgoals=subs, proof=comb, ok=ok, rocq=out)
            if not ok:
                continue
            names = []
            for i, s in enumerate(subs):
                n = self.prove(s, depth + 1, path + [i + 1])
                if n is None:
                    break
                names.append(n)
            if len(names) != len(subs):
                continue
            proof = comb
            for i, n in enumerate(names):
                proof = re.sub(rf"\bsub{i+1}\b", n, proof)
            ok, out = self.check([], self.lemmas, name, stmt, proof)
            self.record(kind="check", stage="assembled", path=path, stmt=stmt, proof=proof, ok=ok, rocq=out)
            if ok:
                self.lemmas.append((name, stmt, proof)); return name
        return None

    def fresh(self):
        self.counter += 1
        return f"{self.task}_L{self.counter}"


def final_file(run, top):
    body = PREAMBLE
    for n, s, p in run.lemmas:
        body += f"\nLemma {n} : {s}.\nProof.\n{p}\nQed.\n"
        if n == top:
            break
    body += f"\nEnd Pilot.\n\nPrint Assumptions {top}.\n"
    return body


def main(tasks):
    summary = {}
    for task in tasks:
        stmt = TASKS[task]
        run = Run(task)
        t0 = time.time()
        # baseline: symbolic automation only
        base = None
        for tac in AUTOMATION:
            ok, _ = run.check([], [], f"{task}_baseline", stmt, tac)
            if ok:
                base = tac; break
        run.record(kind="baseline", stmt=stmt, solved=base is not None, tactic=base)
        top = run.prove(stmt, 0, [])
        result = {"baseline_solved": base is not None, "pilot_solved": top is not None, "model_calls": run.calls,
                  "seconds": round(time.time() - t0, 1)}
        if top:
            f = HERE / "proofs" / f"{task}.v"; f.parent.mkdir(exist_ok=True)
            f.write_text(final_file(run, top))
            with tempfile.TemporaryDirectory() as d:
                g = Path(d) / "Final.v"; g.write_text(f.read_text())
                r = subprocess.run(["coqc", "-q", str(g)], capture_output=True, text=True, timeout=60)
            text = r.stdout + r.stderr
            axioms = re.findall(r"^(\S+) :", text.split("Axioms:")[-1], re.M) if "Axioms:" in text else []
            result.update(final_compiles=r.returncode == 0, forbidden=bool(re.search(r"\b(Admitted|admit|Axiom)\b", f.read_text())),
                          axioms=axioms, closed="Closed under the global context" in text, lemmas_used=len(run.lemmas))
        run.record(kind="result", **result)
        summary[task] = result
        print(task, json.dumps(result), flush=True)
    (HERE / "summary.json").write_text(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main(sys.argv[1:] or list(TASKS))