#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.11"
# dependencies = ["gremlinpython", "python-dotenv"]
# ///
"""
ingest_math_modules.py — Ingest full math metadata into Gremlin

For every .lean file: parses sorry count, theorem/def/structure/inductive counts,
namespace, and dominant math kind. Updates existing module vertices with these
properties, then adds RRC equation vertices linked to their implementing modules.

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

import os
import re
import subprocess
import time
from pathlib import Path
from dotenv import load_dotenv

ROOT     = Path(__file__).resolve().parent.parent
LEAN_DIR = ROOT / "0-Core-Formalism/lean/Semantics/Semantics"
ENV_FILE = ROOT / ".env.gremlin"
load_dotenv(ENV_FILE)

NEON_HOST = "neon-64gb"
CONTAINER = "arxiv-pg"
DB        = "arxiv"

BATCH = 30  # small batches — free tier RU limit

# ── Lean file parser ──────────────────────────────────────────────────────────

MATH_KINDS = {
    "braid":        ["braid","crossing","slug3","ternary","anyon","fusion","topo"],
    "number_theory":["sidon","erdos","prime","diophantine","goormaghtigh","zeckendorf"],
    "geometry":     ["manifold","geodesic","curvature","christoffel","riemannian","spherion"],
    "algebra":      ["quaternion","algebra","ring","field","group","monoid","category"],
    "analysis":     ["convergence","continuity","lipschitz","cauchy","measure","integral"],
    "quantum":      ["yangmills","lattice","hamiltonian","energy","entropy","quantum"],
    "fixedpoint":   ["fix16","q16","fixedpoint","saturate","phasemodulus"],
    "routing":      ["route","cfd","navier","burgers","canal","flow","pressure"],
}

def dominant_kind(text: str) -> str:
    text_l = text.lower()
    scores = {k: sum(text_l.count(w) for w in ws) for k, ws in MATH_KINDS.items()}
    best = max(scores, key=scores.get)
    return best if scores[best] > 0 else "general"

def parse_lean(path: Path) -> dict:
    try:
        text = path.read_text(errors="replace")
    except Exception:
        return {}

    # Module id
    rel = path.relative_to(ROOT / "0-Core-Formalism/lean/Semantics")
    module_id = ".".join(rel.with_suffix("").parts)

    # Namespace
    ns_match = re.search(r"^namespace\s+(\S+)", text, re.MULTILINE)
    namespace = ns_match.group(1) if ns_match else ""

    return {
        "id":          module_id,
        "namespace":   namespace,
        "sorry_count": text.count("sorry"),
        "theorem_count": len(re.findall(r"^theorem\s", text, re.MULTILINE)),
        "def_count":     len(re.findall(r"^def\s",     text, re.MULTILINE)),
        "struct_count":  len(re.findall(r"^structure\s",text, re.MULTILINE)),
        "inductive_count": len(re.findall(r"^inductive\s",text, re.MULTILINE)),
        "line_count":  text.count("\n"),
        "math_kind":   dominant_kind(text),
    }

# ── Gremlin ───────────────────────────────────────────────────────────────────

from gremlin_python.driver import client as gc, serializer

_client = None
def gremlin():
    global _client
    if _client is None:
        _client = gc.Client(
            os.environ["GREMLIN_ENDPOINT"], "g",
            username=os.environ["GREMLIN_USERNAME"],
            password=os.environ["GREMLIN_PASSWORD"],
            message_serializer=serializer.GraphSONSerializersV2d0(),
        )
    return _client

def gq(q: str, b: dict = None) -> list:
    try:
        return gremlin().submitAsync(q, b or {}).result().all().result()
    except Exception as e:
        print(f"  ERR: {str(e)[:100]}")
        return []

def update_module_vertex(m: dict):
    """Add math metadata properties to existing module vertex."""
    gq(
        "g.V().has('module','id',vid)"
        ".property('namespace',ns)"
        ".property('sorry_count',sc)"
        ".property('theorem_count',tc)"
        ".property('def_count',dc)"
        ".property('struct_count',stc)"
        ".property('inductive_count',ic)"
        ".property('line_count',lc)"
        ".property('math_kind',mk)",
        {
            "vid": m["id"],
            "ns":  m["namespace"],
            "sc":  m["sorry_count"],
            "tc":  m["theorem_count"],
            "dc":  m["def_count"],
            "stc": m["struct_count"],
            "ic":  m["inductive_count"],
            "lc":  m["line_count"],
            "mk":  m["math_kind"],
        }
    )

# ── Neon: load RRC equations ──────────────────────────────────────────────────

def neon(sql: str) -> list[list[str]]:
    r = subprocess.run(
        ["ssh", NEON_HOST,
         f"podman exec {CONTAINER} psql -U postgres -d {DB} -t -A -F '|' -c \"{sql}\""],
        capture_output=True, text=True, timeout=60,
    )
    return [ln.split("|") for ln in r.stdout.strip().split("\n") if ln]

def load_rrc_equations() -> list[dict]:
    rows = neon("SELECT equation_id, name FROM rrc_equation_codes8 ORDER BY name")
    return [{"id": r[0], "name": r[1]} for r in rows if len(r) >= 2]

def upsert_equation_vertex(eq: dict):
    gq(
        "g.V().has('equation','id',vid).fold()"
        ".coalesce(unfold(),"
        "  addV('equation').property('id',vid).property('pk',vid).property('name',nm))"
        ".property('name',nm)",
        {"vid": eq["id"], "nm": eq["name"]}
    )

def link_equation_to_modules(eq: dict, module_ids: list[str]):
    """Create 'implements' edges from module → equation."""
    for mid in module_ids:
        gq(
            "g.V().has('module','id',src).as('m')"
            ".V().has('equation','id',dst)"
            ".coalesce("
            "  __.in('implements').where(eq('m')),"
            "  addE('implements').from('m')"
            ")",
            {"src": mid, "dst": eq["id"]}
        )

# ── Module → equation linking heuristic ──────────────────────────────────────

EQ_MODULE_MAP = {
    # equation name substring → module name substrings
    "Chirality":       ["Chirality","Braid","SLUG3","Topo"],
    "BLINK_GATE":      ["SLUG3","Ternary","Blink","Gate"],
    "Christoffel":     ["Christoffel","Geometry","Manifold","Classical"],
    "Stereographic":   ["Spherion","Topo","Geometry","Chart"],
    "Affine_Mapping":  ["Affine","LTSF","Linear","Mapping"],
    "BitFlip":         ["Gradient","Flip","Signal","Energy"],
    "DAG_Force":       ["DAG","Force","Equilibrium","YangMills"],
    "Energy_Function": ["Energy","Entropy","Hamiltonian","Physics"],
    "Energy_Monoton":  ["Energy","Monoton","Convergence","AVMR"],
}

def modules_for_equation(eq_name: str, all_module_ids: list[str]) -> list[str]:
    hits = []
    for prefix, kws in EQ_MODULE_MAP.items():
        if prefix.lower() in eq_name.lower():
            for mid in all_module_ids:
                short = mid.split(".")[-1]
                if any(kw.lower() in short.lower() for kw in kws):
                    hits.append(mid)
    return list(set(hits))

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

def main():
    # ── 1. Parse all .lean files ──
    files = sorted(LEAN_DIR.rglob("*.lean"))
    print(f"Parsing {len(files)} .lean files...")
    modules = []
    for f in files:
        m = parse_lean(f)
        if m:
            modules.append(m)
    print(f"  parsed {len(modules)} modules")

    # Summary
    total_sorry  = sum(m["sorry_count"]   for m in modules)
    total_thm    = sum(m["theorem_count"] for m in modules)
    total_def    = sum(m["def_count"]     for m in modules)
    total_struct = sum(m["struct_count"]  for m in modules)
    kind_counts  = {}
    for m in modules:
        kind_counts[m["math_kind"]] = kind_counts.get(m["math_kind"], 0) + 1
    print(f"  sorries={total_sorry}  theorems={total_thm}  defs={total_def}  structs={total_struct}")
    print(f"  math kinds: {dict(sorted(kind_counts.items(), key=lambda x: -x[1]))}")
    print()

    # ── 2. Update existing module vertices with math metadata ──
    print(f"Updating {len(modules)} module vertices with math metadata...")
    for i, m in enumerate(modules):
        update_module_vertex(m)
        if (i+1) % BATCH == 0:
            print(f"  {i+1}/{len(modules)}", end="\r")
            time.sleep(0.05)
    print(f"\n  done.")

    # ── 3. Load RRC equations from Neon ──
    print("\nLoading RRC equations from Neon...")
    equations = load_rrc_equations()
    print(f"  {len(equations)} equations")

    all_module_ids = [m["id"] for m in modules]

    print(f"Ingesting {len(equations)} equation vertices...")
    for i, eq in enumerate(equations):
        upsert_equation_vertex(eq)
        # Link to implementing modules
        mods = modules_for_equation(eq["name"], all_module_ids)
        if mods:
            link_equation_to_modules(eq, mods[:5])
        if (i+1) % BATCH == 0:
            print(f"  {i+1}/{len(equations)}", end="\r")
            time.sleep(0.05)
    print(f"\n  done.")

    # ── 4. Final counts ──
    v = gq("g.V().count()")
    e = gq("g.E().count()")
    print(f"\nGraph: {v} vertices, {e} edges")

    # ── 5. Top sorry modules ──
    print("\nTop 10 modules by sorry count:")
    results = gq(
        "g.V().hasLabel('module').has('sorry_count', gt(0))"
        ".order().by('sorry_count', decr).limit(10)"
        ".project('module','sorries','kind')"
        ".by('id').by('sorry_count').by('math_kind')"
    )
    if results:
        for r in results:
            print(f"  {r.get('module','?').split('.')[-1]:40s}  sorry={r.get('sorries','?')}  ({r.get('kind','?')})")

    # ── 6. Most connected math kinds ──
    print("\nModules by math_kind:")
    for kind in sorted(kind_counts, key=lambda k: -kind_counts[k])[:6]:
        count = kind_counts[kind]
        print(f"  {kind:20s} {count} modules")

    gremlin().close()
    print("\nIngest complete.")

if __name__ == "__main__":
    main()