#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.11"
# dependencies = ["gremlinpython", "python-dotenv"]
# ///
"""
rrc_math_xref.py — Cross-reference RRC equations → arxiv papers → Lean modules

For each named RRC equation, finds top Jaccard-matching arxiv papers,
then queries Gremlin for Lean modules whose names overlap with paper keywords.
Prints a full cross-reference report.
"""

import os
import subprocess
import sys
from pathlib import Path
from dotenv import load_dotenv

ROOT     = Path(__file__).resolve().parent.parent
ENV_FILE = ROOT / ".env.gremlin"
load_dotenv(ENV_FILE)

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

EQUATIONS = [
    "Chirality_Algebra",
    "BLINK_GATE_Ternary_Clock",
    "Christoffel_Symbols_2D",
    "Stereographic_Chart_Transition",
    "Affine_Mapping_Periodic_Theorem",
    "Affine_Mapping_LTSF_Linear_Layer",
    "BitFlip_Gradient_5D",
    "DAG_Force_Equilibrium",
    "Energy_Function",
    "Energy_Monotonicity_Theorem",
]

# ── Neon ──────────────────────────────────────────────────────────────────────

def neon(sql: str, timeout: int = 90) -> 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=timeout,
    )
    return [ln.split("|") for ln in r.stdout.strip().split("\n") if ln]

def jaccard_top(eq_name: str, k: int = 5) -> list[dict]:
    rows = neon(f"""
WITH eq_codes AS (
  SELECT unnest(codes) AS code, array_length(codes,1) AS eq_len
  FROM rrc_equation_codes8 WHERE name = '{eq_name}'
),
eq_meta AS (SELECT MAX(eq_len) AS eq_len FROM eq_codes),
scored AS (
  SELECT pc.paper_id, COUNT(x.code) AS isect,
    array_length(pc.codes,1) + m.eq_len - COUNT(x.code) AS union_sz
  FROM arxiv_paper_codes8 pc JOIN eq_meta m ON true JOIN eq_codes x ON x.code = ANY(pc.codes)
  GROUP BY pc.paper_id, pc.codes, m.eq_len
)
SELECT s.paper_id, ROUND(s.isect::numeric/s.union_sz,3), ap.title, ap.categories
FROM scored s JOIN arxiv_papers ap ON s.paper_id = ap.paper_id
ORDER BY 2 DESC LIMIT {k}
""".replace("\n", " "))
    return [{"id": r[0], "j": r[1], "title": r[2], "cat": r[3]} for r in rows if len(r) >= 3]

# ── 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 gremlin_q(q: str, b: dict = None) -> list:
    try:
        return gremlin().submitAsync(q, b or {}).result().all().result()
    except Exception as e:
        return []

def modules_for_keywords(keywords: list[str]) -> list[str]:
    """Find internal Lean modules whose id contains any keyword."""
    all_ids = gremlin_q(
        "g.V().hasLabel('module').has('kind','internal').values('id')"
    )
    hits = []
    for mid in all_ids:
        for kw in keywords:
            if kw.lower() in mid.lower():
                hits.append(mid)
                break
    return sorted(set(hits))

def module_in_degree(module_id: str) -> int:
    r = gremlin_q(
        "g.V().has('module','id',mid).in('imports').count()",
        {"mid": module_id}
    )
    return r[0] if r else 0

# ── Keyword extraction ────────────────────────────────────────────────────────

STOP = {"and","the","of","a","in","on","for","with","to","from","by","an","is","are","via"}

def title_keywords(title: str) -> list[str]:
    words = title.lower().replace("-"," ").replace("$","").split()
    return [w.strip(".,()[]{}") for w in words if len(w) > 3 and w not in STOP]

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

def main():
    print("Loading all internal module IDs from Gremlin...")
    all_modules = gremlin_q(
        "g.V().hasLabel('module').has('kind','internal').values('id')"
    )
    print(f"  {len(all_modules)} internal modules\n")

    for eq in EQUATIONS:
        print(f"{'═'*70}")
        print(f"  {eq}")
        print(f"{'─'*70}")

        papers = jaccard_top(eq, k=5)
        if not papers:
            print("  (no Neon matches)\n")
            continue

        # Collect keywords from top paper titles
        all_kws = []
        for p in papers:
            all_kws.extend(title_keywords(p["title"]))

        # Find Lean modules matching keywords
        matching_mods = []
        for mid in all_modules:
            short = mid.split(".")[-1].lower()
            for kw in all_kws:
                if kw in short or short in kw:
                    matching_mods.append(mid)
                    break
        matching_mods = sorted(set(matching_mods))

        print(f"  arxiv matches:")
        for p in papers:
            print(f"    [{p['j']}] {p['id']:15s} {p['title'][:55]}  ({p['cat']})")

        print(f"  Lean modules ({len(matching_mods)}):")
        for m in matching_mods[:8]:
            deg = module_in_degree(m)
            print(f"    {m:55s} ← {deg} importers")

        print()

    gremlin().close()

if __name__ == "__main__":
    main()