"""Per-theorem SAT instance sizes for the bv_decide proofs.

Re-elaborates each module with `trace.Meta.Tactic.bv` enabled and `--json` output: every trace
block is one JSON message carrying its source position, so each "AIG has N nodes" value (size of
the bit-blasted And-Inverter Graph handed to CaDiCaL) is attributed to the enclosing theorem
exactly. `bv_decide` calls closed by preprocessing alone emit no AIG. Requires a prior
`lake build` (imports must be compiled).

Run (from formal/):  ../.venv/bin/python scripts/aig_sizes.py
"""
from __future__ import annotations
import json, os, re, subprocess, sys, time
from pathlib import Path

HERE = Path(__file__).resolve().parent.parent
MODULES = ["Sha256Words", "Sha256Step", "Schedule", "LocalCollision"]

def theorem_at(src_lines, lineno):
    name = None
    for i in range(lineno - 1, -1, -1):
        m = re.match(r"\s*(?:theorem|lemma|def)\s+([\w.']+)", src_lines[i])
        if m:
            return m.group(1)
    return name

def main():
    env = dict(os.environ)
    env["PATH"] = str(Path.home() / ".elan" / "bin") + os.pathsep + env.get("PATH", "")
    rows = []
    for mod in MODULES:
        path = HERE / "Wp34" / f"{mod}.lean"
        src = path.read_text().splitlines()
        t0 = time.time()
        r = subprocess.run(["lake", "env", "lean", "--json", "-D", "trace.Meta.Tactic.bv=true", str(path)],
                           cwd=HERE, env=env, capture_output=True, text=True)
        dt = time.time() - t0
        for line in r.stdout.splitlines():
            try:
                msg = json.loads(line)
            except json.JSONDecodeError:
                continue
            for m in re.finditer(r"AIG has (\d+) nodes", msg.get("data", "")):
                rows.append((mod, theorem_at(src, msg["pos"]["line"]), int(m.group(1))))
        print(f"# {mod}: re-elaborated with trace in {dt:.1f} s, exit {r.returncode}", file=sys.stderr)
    print("module\ttheorem\taig_nodes")
    for mod, name, n in rows:
        print(f"{mod}\t{name}\t{n}")

if __name__ == "__main__":
    main()