Research-Stack/scripts/load_dependency_graph.py

#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "gremlinpython",
#   "python-dotenv",
# ]
# ///
"""
load_dependency_graph.py — Full cross-layer dependency graph for RRC.

Creates 6 vertex types and 8 edge types in the mathblob Gremlin database:

  Vertices:  module, theorem, rrc_equation, receipt, shim, hardware_probe
  Edges:     imports, contains, proves, certifies, stamps, extracts, probes, depends_on

Phases:
  1. Lean module + theorem parsing (all 840 .lean files)
  2. RRC equation extraction (Corpus250.lean)
  3. Receipt extraction (AVMIsa.Emit + receipt JSONs)
  4. Shim extraction (4-Infrastructure/shim/)
  5. Hardware probe extraction (4-Infrastructure/hardware/)
  6. Gremlin loading (mathblob)
  7. Verification queries

Run:  uv run scripts/load_dependency_graph.py
"""

import concurrent.futures
import json
import os
import re
import sys
import time
from collections import defaultdict
from pathlib import Path
from dotenv import load_dotenv
from gremlin_python.driver import client as gremlin_client, serializer

# ── Config ────────────────────────────────────────────────────────────────────

ROOT     = Path(__file__).parent.parent
LEAN_DIR = ROOT / "0-Core-Formalism/lean/Semantics/Semantics"
SHIM_DIR = ROOT / "4-Infrastructure/shim"
HW_DIR   = ROOT / "4-Infrastructure/hardware"
RCPT_DIR = ROOT / "shared-data/data/stack_solidification"
ENV_FILE = ROOT / ".env.gremlin"

load_dotenv(ENV_FILE)

ENDPOINT = os.environ["GREMLIN_ENDPOINT"]
USERNAME = os.environ["GREMLIN_USERNAME"]
PASSWORD = os.environ["GREMLIN_PASSWORD"]



# ── Regex patterns ────────────────────────────────────────────────────────────

RE_IMPORT  = re.compile(r'^import\s+(\S+)')
RE_THEOREM = re.compile(
    r'^\s*(private\s+|protected\s+)?(theorem|lemma|def)\s+(\w+)'
)
RE_CAPITALIZED_IDENT = re.compile(r'\b([A-Z][A-Za-z0-9_.]*)\b')
RE_DOT_REF = re.compile(r'\b(\w+\.\w+)\b')

# Shim patterns: calls to Lean from Python
RE_EVAL      = re.compile(r'#eval\s+\S+')
RE_LAKE_EXE  = re.compile(r'lake\s+(?:build|exe|run)\s+(\S+)')
RE_LAKE_BUILD = re.compile(r'lake build (\S+)')
RE_EVAL_CALL = re.compile(r'(leanBuildReceipt|emitCorpus|emitRrc|evalSparsePoly|verifyCertificate)\b')

# ── Phase 1: Module + Theorem parsing ────────────────────────────────────────

def parse_lean_file(path: Path) -> dict:
    """Return {id, imports, theorems: [{name, kind, line, body_lines}]}."""
    rel = path.relative_to(ROOT / "0-Core-Formalism/lean/Semantics")
    module_id = ".".join(rel.with_suffix("").parts)

    result = {
        "id": module_id,
        "label": "module",
        "kind": "internal",
        "file_path": str(rel),
        "line_count": 0,
        "imports": [],
        "theorems": [],
    }

    try:
        text = path.read_text(errors="replace")
        lines = text.splitlines()
        result["line_count"] = len(lines)
    except Exception as e:
        print(f"  WARN: could not read {path.name}: {e}")
        return result

    # Parse imports (first contiguous block)
    in_imports = True
    for i, line in enumerate(lines):
        if in_imports:
            m = RE_IMPORT.match(line.strip())
            if m:
                result["imports"].append(m.group(1))
                continue
            stripped = line.strip()
            if stripped == "" or stripped.startswith("--"):
                continue
            # First non-import, non-blank, non-comment line ends import block
            # But allow comment-only files
            if not stripped.startswith("--") and stripped:
                in_imports = False

    # Parse theorem/lemma/def declarations
    i = 0
    while i < len(lines):
        line = lines[i]
        m = RE_THEOREM.match(line)
        if m:
            thm_name = m.group(3)
            thm_kind = m.group(2)  # theorem, lemma, or def
            thm_line = i + 1  # 1-indexed

            # Collect body lines until next top-level declaration
            body_lines = []
            depth = 0
            in_body = False
            j = i + 1
            # Find the := or : or where
            decl_line = line
            while j < len(lines) and not lines[j].lstrip().startswith(("theorem ", "lemma ", "def ")):
                # Track brace depth for where clauses
                for ch in lines[j]:
                    if ch == '{': depth += 1; in_body = True
                    elif ch == '}': depth -= 1
                if depth > 0 or in_body:
                    body_lines.append(lines[j])
                elif ':= by' in lines[j] or ':=' in lines[j]:
                    in_body = True
                    body_lines.append(lines[j])
                elif (lines[j].strip().startswith(":= by") or
                      lines[j].strip().startswith(":=")):
                    in_body = True
                    body_lines.append(lines[j])
                j += 1

            result["theorems"].append({
                "name": thm_name,
                "kind": thm_kind,
                "line": thm_line,
                "decl_line": decl_line,
                "body_lines": body_lines,
            })
            i = j
        else:
            i += 1

    return result


# ── Phase 1b: Theorem dependency extraction ──────────────────────────────────

def extract_theorem_deps(module_id: str, theorems: list[dict],
                         all_theorems: dict[str, str]) -> list[tuple]:
    """
    Return [(src_thm_id, dst_thm_id, 'proves')] for each theorem's body
    referencing another known theorem.

    all_theorems: { theorem_name: module_id } — global index of all theorems.
    """
    edges = []
    for thm in theorems:
        src_id = f"{module_id}.{thm['name']}"
        body_text = "\n".join(thm["body_lines"])

        # Find all potential references
        refs = set()
        for m in RE_CAPITALIZED_IDENT.finditer(body_text):
            name = m.group(1)
            if name in all_theorems:
                refs.add(name)

        for m in RE_DOT_REF.finditer(body_text):
            # module.TheoremName pattern
            parts = m.group(1).split(".")
            if len(parts) == 2 and parts[1] in all_theorems:
                refs.add(parts[1])

        # Filter out self-references and local vars
        for ref in refs:
            if ref == thm['name']:
                continue
            # Skip local 'h_' hypothesis references
            if ref.startswith('h') and len(ref) > 1 and ref[1].islower() and len(ref) < 20:
                continue
            dst_module = all_theorems[ref]
            dst_id = f"{dst_module}.{ref}"
            if dst_id != src_id:
                edges.append((src_id, dst_id, "proves"))

    return edges


# ── Phase 1 main: collect all module + theorem data ──────────────────────────

def collect_lean_graph() -> tuple[dict, list, dict]:
    """
    Walk all .lean files, return:
      vertices: { vertex_id: {id, label, kind, ...} }
      edges:    [(src, dst, label), ...]
      all_theorems: { theorem_name: module_id }
    """
    files = sorted(LEAN_DIR.rglob("*.lean"))
    print(f"\n── Phase 1: Scanning {len(files)} .lean files ──")

    vertices = {}
    edges = []
    parsed_modules = []
    all_theorems = {}

    for f in files:
        mod = parse_lean_file(f)
        mid = mod["id"]
        vertices[mid] = {
            "id": mid,
            "label": "module",
            "kind": mod["kind"],
            "file_path": mod["file_path"],
            "line_count": mod["line_count"],
            "theorem_count": len(mod["theorems"]),
        }

        for imp in mod["imports"]:
            if imp not in vertices:
                kind = "external" if not imp.startswith("Semantics.") else "internal"
                vertices[imp] = {"id": imp, "label": "module", "kind": kind}
            edges.append((mid, imp, "imports"))

        # Register theorems
        for thm in mod["theorems"]:
            thm_id = f"{mid}.{thm['name']}"
            all_theorems[thm['name']] = mid
            vertices[thm_id] = {
                "id": thm_id,
                "label": "theorem",
                "module_id": mid,
                "kind": thm["kind"],
                "line": thm["line"],
            }
            edges.append((mid, thm_id, "contains"))

        parsed_modules.append((mid, mod["theorems"]))

    # Second pass: extract theorem dependencies (need global theorem index)
    print(f"  Extracting theorem dependencies ({len(all_theorems)} theorems)...")
    dep_count = 0
    for mid, theorems in parsed_modules:
        deps = extract_theorem_deps(mid, theorems, all_theorems)
        edges.extend(deps)
        dep_count += len(deps)

    print(f"  {len(vertices)} vertices ({len([v for v in vertices.values() if v['label']=='module'])} modules, "
          f"{len([v for v in vertices.values() if v['label']=='theorem'])} theorems), "
          f"{len(edges)} edges ({dep_count} proof deps)")
    return vertices, edges, all_theorems


# ── Phase 2: RRC equation vertices ──────────────────────────────────────────

def parse_corpus250() -> list[dict]:
    """Extract equation data from Corpus250.lean."""
    path = LEAN_DIR / "RRC/Corpus250.lean"
    if not path.exists():
        print("  Corpus250.lean not found, skipping")
        return []
    print("\n── Phase 2: Parsing RRC equations ──")
    text = path.read_text(errors="replace")

    # Split on `},\n  { equationId` — each boundary is consumed, so fragments
    # after the first start with ` := "..."`. Reconstruct with `{ equationId`.
    fragments = text.split("},\n  { equationId")

    equations = []
    for i, frag in enumerate(fragments):
        # First fragment: preamble before first "{ equationId", or the row
        # Subsequent fragments: start with ` := "..."`, missing the opening
        if i == 0:
            idx = frag.find("{ equationId")
            if idx < 0:
                continue
            row_text = frag[idx:]
        else:
            row_text = "{ equationId" + frag

        eq = {}
        m = re.search(r'equationId\s*:=\s*"([^"]+)"', row_text)
        if m: eq["equationId"] = m.group(1)
        m = re.search(r'name\s*:=\s*"([^"]+)"', row_text)
        if m: eq["name"] = m.group(1)
        m = re.search(r'shape\s*:=\s*\.(\w+)', row_text)
        if m: eq["shape"] = m.group(1)
        m = re.search(r'status\s*:=\s*\.(\w+)', row_text)
        if m: eq["status"] = m.group(1)
        m = re.search(r'fourcc\s*:=\s*"([^"]+)"', row_text)
        if m: eq["fourcc"] = m.group(1)
        m = re.search(r'rrcKind\s*:=\s*"([^"]+)"', row_text)
        if m: eq["rrc_kind"] = m.group(1)
        m = re.search(r'weakAxesCnt\s*:=\s*(\d+)', row_text)
        if m: eq["weak_axes_cnt"] = int(m.group(1))
        m = re.search(r'arxivPaperId\s*:=\s*some\s*"([^"]+)"', row_text)
        if m: eq["arxiv_paper_id"] = m.group(1)
        m = re.search(r'templateKey\s*:=\s*"([^"]+)"', row_text)
        if m: eq["template_key"] = m.group(1)

        if "equationId" in eq and "name" in eq:
            equations.append(eq)

    print(f"  Extracted {len(equations)} equations from Corpus250.lean")
    return equations


def add_rrc_vertices_edges(equations: list[dict], all_theorems: dict[str, str]
                           ) -> tuple[dict, list]:
    """Add RRC equation vertices and certification edges."""
    vertices = {}
    edges = []

    for eq in equations:
        eq_id = eq["equationId"]
        vertices[eq_id] = {
            "id": eq_id,
            "label": "rrc_equation",
            "name": eq.get("name", ""),
            "shape": eq.get("shape", ""),
            "status": eq.get("status", ""),
            "arxiv_paper_id": eq.get("arxivPaperId", ""),
        }

        # Link to certifying theorem if name matches a known theorem
        eq_name = eq.get("name", "")
        for thm_name in all_theorems:
            if thm_name.lower() in eq_name.lower() or eq_name.lower() in thm_name.lower():
                dst_module = all_theorems[thm_name]
                dst_id = f"{dst_module}.{thm_name}"
                edges.append((eq_id, dst_id, "certifies"))

    print(f"  {len(vertices)} RRC equations, {len(edges)} certification edges")
    return vertices, edges


# ── Phase 3: Receipt vertices ────────────────────────────────────────────────

def parse_avm_receipts() -> list[dict]:
    """Extract receipt data from AVMIsa.Emit.lean."""
    path = LEAN_DIR / "AVMIsa/Emit.lean"
    receipts = []
    if not path.exists():
        return receipts
    text = path.read_text(errors="replace")

    # Canary receipts: canaryReceipt "targetId" prog expected
    for m in re.finditer(r'canaryReceipt\s+"([^"]+)"', text):
        receipts.append({
            "id": m.group(1),
            "kind": "leanBuild",
            "targetId": m.group(1),
            "valid": True,
            "authority": "avm",
        })

    # Bundle receipt
    if 'leanBuildReceipt "avm.rrc_corpus250.bundle"' in text:
        receipts.append({
            "id": "avm.rrc_corpus250.bundle",
            "kind": "leanBuild",
            "targetId": "avm.rrc_corpus250.bundle",
            "valid": True,
            "authority": "avm",
        })

    return receipts


def parse_receipt_json_file(f: Path) -> list[dict]:
    """Parse a single receipt JSON file, return list of receipt dicts."""
    try:
        data = json.loads(f.read_text())
    except (json.JSONDecodeError, Exception) as e:
        print(f"  WARN: could not parse {f.name}: {e}")
        return []

    if isinstance(data, list):
        items = data
    elif isinstance(data, dict):
        items = [data]
    else:
        return []

    results = []
    for item in items:
        receipt_id = item.get("receipt_hash") or item.get("id") or f.stem
        results.append({
            "id": receipt_id,
            "kind": item.get("schema", "stack_receipt"),
            "targetId": item.get("target_id") or item.get("claim_boundary", ""),
            "valid": item.get("action_result", {}).get("status") == "ok" if "action_result" in item else True,
            "authority": item.get("authority", "system"),
            "file": f.name,
        })
    return results


def parse_receipt_jsons() -> list[dict]:
    """Parse all receipt JSON files from stack_solidification/."""
    receipts = []
    if not RCPT_DIR.exists():
        return receipts

    for f in sorted(RCPT_DIR.glob("*.json")):
        receipts.extend(parse_receipt_json_file(f))

    return receipts


def add_receipt_vertices_edges(receipts: list[dict]) -> tuple[dict, list]:
    """Add receipt vertices and stamps edges to equations."""
    vertices = {}
    edges = []

    for r in receipts:
        rid = r["id"]
        vertices[rid] = {
            "id": rid,
            "label": "receipt",
            "kind": r["kind"],
            "target_id": r["targetId"],
            "valid": r["valid"],
            "authority": r.get("authority", ""),
            "file": r.get("file", ""),
        }

        # Link receipt to target (if it's an RRC equation ID pattern)
        tid = r["targetId"]
        if tid.startswith("rrc_eq_"):
            edges.append((rid, tid, "stamps"))
        # Bundle receipt stamps the whole corpus
        if tid == "avm.rrc_corpus250.bundle":
            edges.append((rid, "rrc.corpus250", "stamps"))

    return vertices, edges


# ── Phase 4: Shim vertices ───────────────────────────────────────────────────

def walk_shims() -> list[dict]:
    """Walk 4-Infrastructure/shim/ for Python scripts."""
    shims = []
    if not SHIM_DIR.exists():
        return shims

    for f in sorted(SHIM_DIR.glob("*.py")):
        try:
            text = f.read_text(errors="replace")
        except Exception:
            continue

        # Detect references to Lean theorems/modules
        lean_refs = set()
        for m in RE_EVAL_CALL.finditer(text):
            lean_refs.add(m.group(1))

        # Detect #eval statements in docstrings/comments referencing Lean
        for line in text.splitlines():
            if '#eval' in line or 'lake build' in line or 'lake exe' in line:
                lean_refs.add(line.strip()[:80])

        shims.append({
            "id": f"shim://{f.relative_to(ROOT)}",
            "path": str(f.relative_to(ROOT)),
            "type": "python",
            "line_count": len(text.splitlines()),
            "lean_refs": list(lean_refs),
        })

    print(f"  Found {len(shims)} shim files")
    return shims


def walk_hardware_probes() -> list[dict]:
    """Walk 4-Infrastructure/hardware/ for probe scripts."""
    probes = []
    if not HW_DIR.exists():
        return probes

    for f in sorted(HW_DIR.glob("*probe*.py")) + sorted(HW_DIR.glob("*probe*.sh")):
        try:
            text = f.read_text(errors="replace")
        except Exception:
            continue

        probes.append({
            "id": f"probe://{f.relative_to(ROOT)}",
            "path": str(f.relative_to(ROOT)),
            "type": f.suffix.lstrip("."),
            "line_count": len(text.splitlines()),
        })

    print(f"  Found {len(probes)} hardware probes")
    return probes


def add_shim_vertices_edges(shims: list[dict], all_theorems: dict[str, str]
                             ) -> tuple[dict, list]:
    """Add shim vertices and extraction edges to theorems."""
    vertices = {}
    edges = []

    for sh in shims:
        sid = sh["id"]
        vertices[sid] = {
            "id": sid,
            "label": "shim",
            "path": sh["path"],
            "type": sh["type"],
            "line_count": sh["line_count"],
        }

        # Link shim to known theorems via #eval references
        for ref_text in sh["lean_refs"]:
            for thm_name in all_theorems:
                if thm_name.lower() in ref_text.lower():
                    dst_module = all_theorems[thm_name]
                    dst_id = f"{dst_module}.{thm_name}"
                    edges.append((sid, dst_id, "extracts"))

    return vertices, edges


def add_hw_vertices_edges(probes: list[dict], shim_ids: set) -> tuple[dict, list]:
    """Add hardware probe vertices and probes edges to shims."""
    vertices = {}
    edges = []

    for p in probes:
        pid = p["id"]
        vertices[pid] = {
            "id": pid,
            "label": "hardware_probe",
            "path": p["path"],
            "type": p["type"],
            "line_count": p["line_count"],
        }

        # Link probe to shims via path similarity
        probe_stem = Path(p["path"]).stem.lower().replace("_probe", "").replace("probe_", "")
        for sid in shim_ids:
            shim_stem = Path(sid.replace("shim://", "")).stem.lower()
            if probe_stem in shim_stem or shim_stem in probe_stem:
                edges.append((pid, sid, "probes"))

    return vertices, edges


# ── Phase 6: Gremlin loading ─────────────────────────────────────────────────

def make_client():
    return gremlin_client.Client(
        ENDPOINT, "g",
        username=USERNAME,
        password=PASSWORD,
        message_serializer=serializer.GraphSONSerializersV2d0(),
    )


GREMLIN_TIMEOUT = 30  # seconds per query

def submit(c, query: str, bindings: dict = None):
    try:
        cb = c.submitAsync(query, bindings or {})
        return cb.result(timeout=GREMLIN_TIMEOUT).all().result()
    except concurrent.futures.TimeoutError:
        print(f"  TIMEOUT ({GREMLIN_TIMEOUT}s): {query[:120]}")
        return None
    except Exception as e:
        print(f"  ERR: {e!s:.200}")
        return None


# def _gremlin_val(val) -> str:
#     """Format a Python value for Gremlin query embedding."""
#     if isinstance(val, bool):
#         return str(val).lower()
#     elif isinstance(val, int):
#         return str(val)
#     elif isinstance(val, str):
#         safe = val.replace("'", "\\'").replace('"', '\\"')
#         return f"'{safe}'"
#     else:
#         return f"'{str(val)}'"


def upsert_vertex(c, v: dict):
    """Add or update a typed vertex — single Gremlin query with bindings.

    Uses the exact pattern from load_module_graph.py (verified working):
      fold().coalesce(unfold(), addV()).property() ...
    """
    label = v["label"]
    vid = _safe_id(v["id"])
    pk = vid  # partition key = sanitized id

    # Build property key-value pairs for bindings
    props = {k: val for k, val in v.items() if k not in ("id", "label", "pk")}

    # Create property chain for the update phase (after coalesce)
    update_props = ""
    for k in props:
        update_props += f".property('{k}',{k})"

    # Create property chain for the create phase (inside addV)
    create_props = ".property('id',idv).property('pk',pkv)"
    for k in props:
        create_props += f".property('{k}',{k})"

    q = (f"g.V().has('{label}','id',idv).fold()"
         f".coalesce("
         f"  unfold(),"
         f"  addV('{label}'){create_props}"
         f"){update_props}")

    b = {"idv": vid, "pkv": pk, **props}
    submit(c, q, b)


def _safe_id(s: str) -> str:
    """Sanitize a Gremlin element ID — Cosmos DB rejects '/'."""
    return s.replace("/", ".").replace("\\", ".").replace("|", ".").replace(" ", "_")

def _esc(s: str) -> str:
    """Escape for Gremlin single-quoted strings."""
    return s.replace("'", "\\'").replace('"', '\\"')


def load_to_gremlin(all_vertices: dict[str, dict], all_edges: list[tuple]):
    """Load all vertices and edges into Gremlin.

    Uses idempotent upsert patterns (verified working in load_module_graph.py):
      - Vertices: fold().coalesce(unfold(), addV()).property()...
      - Edges:    coalesce(select('s').outE(lbl).where(inV().as('d')), addE(lbl).from('s').to('d'))

    No drop() calls — Cosmos DB Gremlin doesn't support them reliably.
    Old edges from prior runs are NOT cleared (accepting this limitation).
    """
    print(f"\n── Phase 6: Loading to Gremlin ({ENDPOINT}) ──")
    c = make_client()

    count = submit(c, "g.V().count()")
    print(f"  Current vertex count: {count}")

    # Load vertices — sequential, one at a time (free tier RU budget)
    edges_labels = {'contains', 'proves', 'certifies', 'stamps', 'extracts', 'probes'}
    # Separate module vertices (already exist) from new dependency vertices
    module_vertices = [v for v in all_vertices.values() if v['label'] == 'module']
    dep_vertices = [v for v in all_vertices.values() if v['label'] != 'module']

    print(f"  Loading {len(module_vertices)} module vertices (upsert)...")
    for i, v in enumerate(module_vertices):
        try:
            upsert_vertex(c, v)
        except Exception as e:
            print(f"    [{v['id']}] upsert failed: {e}")
        if (i + 1) % 100 == 0:
            print(f"    modules {i+1}/{len(module_vertices)}")

    print(f"  Loading {len(dep_vertices)} dependency vertices...")
    for i, v in enumerate(dep_vertices):
        try:
            upsert_vertex(c, v)
        except Exception as e:
            print(f"    [{v['id']}] upsert failed: {e}")
        if (i + 1) % 100 == 0:
            print(f"    deps {i+1}/{len(dep_vertices)}")

    v_total = len(all_vertices)
    print(f"  vertices done ({v_total}).")

    # Load edges — parallel with one Gremlin client per worker thread
    total_e = len(all_edges)
    print(f"  Loading {total_e} edges (4 workers)...")

    from concurrent.futures import ThreadPoolExecutor
    import threading

    thread_local = threading.local()

    def _init_worker():
        thread_local.client = make_client()

    def _send_edge(args):
        src, dst, lbl = args
        cl = thread_local.client
        src_safe = _safe_id(src)
        dst_safe = _safe_id(dst)
        src_esc = _esc(src_safe)
        dst_esc = _esc(dst_safe)
        lbl_esc = _esc(lbl)
        q = (f"g.V().has('id','{src_esc}').as('s')"
             f".V().has('id','{dst_esc}').as('d')"
             f".coalesce("
             f"  select('s').outE('{lbl_esc}').where(inV().as('d')),"
             f"  addE('{lbl_esc}').from('s').to('d')"
             f")")
        try:
            submit(cl, q)
            return True
        except Exception:
            return False

    failed_edges = 0
    with ThreadPoolExecutor(max_workers=4, initializer=_init_worker) as pool:
        for i in range(0, total_e, 100):
            batch = all_edges[i:i+100]
            results = pool.map(_send_edge, batch)
            failed_edges += sum(1 for r in results if not r)
            done = min(i + 100, total_e)
            if done % 500 == 0 or done == total_e:
                print(f"    edges {done}/{total_e}")

    print(f"  edges done ({total_e}, {failed_edges} failed).")

    # Final counts
    time.sleep(1)
    v_count = submit(c, "g.V().count()")
    e_count = submit(c, "g.E().count()")
    type_counts = submit(c,
        "g.V().groupCount().by('label').unfold()"
        ".project('type','count').by(keys).by(values)"
    )
    print(f"\n  Graph loaded: {v_count} vertices, {e_count} edges")
    if type_counts:
        for tc in type_counts:
            print(f"    {tc.get('type','?'):20s}: {tc.get('count',0)}")

    c.close()


# ── Phase 7: Verification ────────────────────────────────────────────────────

def run_verification_queries():
    """Run a comprehensive set of verification queries against the graph."""
    print(f"\n── Phase 7: Verification queries ──")
    c = make_client()

    def q(query, b=None):
        return c.submitAsync(query, b or {}).result().all().result()

    queries = [
        ("Q1: Total vertices by type",
         "g.V().groupCount().by('label').unfold()"
         ".project('type','count').by(keys).by(values)"),
        ("Q2: Total edges by type",
         "g.E().groupCount().by('label').unfold()"
         ".project('type','count').by(keys).by(values)"),
        ("Q3: Top 10 most-imported modules",
         "g.V().hasLabel('module').has('kind','internal')"
         ".order().by(__.in('imports').count(), decr).limit(10)"
         ".project('module','imported_by','theorems')"
         ".by('id').by(__.in('imports').count()).by('theorem_count')"),
        ("Q4: Modules with most theorems",
         "g.V().hasLabel('module').has('kind','internal')"
         ".order().by('theorem_count', decr).limit(10)"
         ".project('module','theorems')"
         ".by('id').by('theorem_count')"),
        ("Q5: Most-referenced theorems (most 'proves' in-edges)",
         "g.V().hasLabel('theorem')"
         ".order().by(__.in('proves').count(), decr).limit(10)"
         ".project('theorem','referenced_by')"
         ".by('id').by(__.in('proves').count())"),
        ("Q6: RRC equations count by status",
         "g.V().hasLabel('rrc_equation')"
         ".groupCount().by('status').unfold()"
         ".project('status','count').by(keys).by(values)"),
        ("Q7: Receipts by kind",
         "g.V().hasLabel('receipt')"
         ".groupCount().by('kind').unfold()"
         ".project('kind','count').by(keys).by(values)"),
        ("Q8: Shim → Lean theorem extraction edges (top 10 shims by refs)",
         "g.V().hasLabel('shim')"
         ".order().by(__.out('extracts').count(), decr).limit(10)"
         ".project('shim','theorems_extracted')"
         ".by('id').by(__.out('extracts').count())"),
        ("Q9: Dependency chains (modules with both imports AND theorem deps)",
         "g.V().hasLabel('module').has('kind','internal')"
         ".where(__.out('imports').count().is(gt(0)))"
         ".where(__.out('contains').count().is(gt(0)))"
         ".limit(5).values('id')"),
        ("Q10: Full chain: theorem → proves → theorem → contains → module",
         "g.V().hasLabel('theorem').limit(3)"
         ".project('theorem','proves','module')"
         ".by('id')"
         ".by(__.out('proves').limit(3).values('id').fold())"
         ".by(__.in('contains').values('id'))"),
    ]

    for name, gremlin_q in queries:
        print(f"\n  {name}")
        print(f"  {'─' * (len(name) + 2)}")
        try:
            results = q(gremlin_q)
            if results:
                for r in results[:8]:
                    if isinstance(r, dict):
                        parts = [f"{k}={v}" for k, v in r.items()]
                        print(f"    {', '.join(parts)}")
                    else:
                        print(f"    {r}")
            else:
                print("    (empty)")
        except Exception as exc:
            print(f"    ERR: {exc!s:.120}")

    c.close()


# ── Main ─────────────────────────────────────────────────────────────────────

def main():
    print("═" * 60)
    print("  RRC Full Dependency Graph Loader")
    print("═" * 60)

    # Phase 1: Lean modules + theorems
    mod_vertices, mod_edges, all_theorems = collect_lean_graph()

    # Phase 2: RRC equations
    equations = parse_corpus250()
    rrc_vertices, rrc_edges = add_rrc_vertices_edges(equations, all_theorems)

    # Phase 3: Receipts
    avm_receipts = parse_avm_receipts()
    json_receipts = parse_receipt_jsons()
    all_receipts = avm_receipts + json_receipts
    print(f"\n── Phase 3: Receipts ──")
    print(f"  {len(avm_receipts)} AVM receipts, {len(json_receipts)} JSON receipts")
    rcpt_vertices, rcpt_edges = add_receipt_vertices_edges(all_receipts)

    # Phase 4: Shims
    shims = walk_shims()
    print(f"\n── Phase 4: Shims → theorems ──")
    shim_vertices, shim_edges = add_shim_vertices_edges(shims, all_theorems)

    # Phase 5: Hardware probes
    probes = walk_hardware_probes()
    print(f"\n── Phase 5: Hardware probes → shims ──")
    shim_ids = set(shim_vertices.keys())
    hw_vertices, hw_edges = add_hw_vertices_edges(probes, shim_ids)

    # Merge all vertices and edges
    all_vertices = {}
    all_vertices.update(mod_vertices)
    all_vertices.update(rrc_vertices)
    all_vertices.update(rcpt_vertices)
    all_vertices.update(shim_vertices)
    all_vertices.update(hw_vertices)

    all_edges = []
    all_edges.extend(mod_edges)
    all_edges.extend(rrc_edges)
    all_edges.extend(rcpt_edges)
    all_edges.extend(shim_edges)
    all_edges.extend(hw_edges)

    print(f"\n{'═' * 60}")
    print(f"  Total: {len(all_vertices)} vertices, {len(all_edges)} edges")
    print(f"    modules:       {len([v for v in all_vertices.values() if v['label']=='module'])}")
    print(f"    theorems:      {len([v for v in all_vertices.values() if v['label']=='theorem'])}")
    print(f"    rrc_equations: {len([v for v in all_vertices.values() if v['label']=='rrc_equation'])}")
    print(f"    receipts:      {len([v for v in all_vertices.values() if v['label']=='receipt'])}")
    print(f"    shims:         {len([v for v in all_vertices.values() if v['label']=='shim'])}")
    print(f"    hw_probes:     {len([v for v in all_vertices.values() if v['label']=='hardware_probe'])}")
    print(f"    imports:       {sum(1 for e in all_edges if e[2]=='imports')}")
    print(f"    contains:      {sum(1 for e in all_edges if e[2]=='contains')}")
    print(f"    proves:        {sum(1 for e in all_edges if e[2]=='proves')}")
    print(f"    certifies:     {sum(1 for e in all_edges if e[2]=='certifies')}")
    print(f"    stamps:        {sum(1 for e in all_edges if e[2]=='stamps')}")
    print(f"    extracts:      {sum(1 for e in all_edges if e[2]=='extracts')}")
    print(f"    probes:        {sum(1 for e in all_edges if e[2]=='probes')}")
    print(f"{'═' * 60}")

    # Phase 6: Load to Gremlin
    load_to_gremlin(all_vertices, all_edges)

    # Phase 7: Verify
    run_verification_queries()

    print("\nDone.")


if __name__ == "__main__":
    main()