#!/usr/bin/env python3
"""
rrc_domain_manifold_graph.py — Build expanding manifold graph across all domains.

Connects all gathered math domains into an expanding manifold graph.
Edges are: shared arxiv papers, keyword overlap, and RRC route connections.

Output: shared-data/data/domain_manifold_graph_v1.json
"""
from __future__ import annotations

import json
import subprocess
import sys
from collections import defaultdict
from pathlib import Path

ROOT = Path("shared-data/data")
OUT_PATH = ROOT / "domain_manifold_graph_v1.json"
NEON_HOST = "neon-64gb"
CONTAINER = "arxiv-pg"
DB = "arxiv"

KERNEL_SOURCES = {
    "diophantine": {
        "file": ROOT / "diophantine_kernel_v1.json",
        "label": "Diophantine / Number Theory",
        "dimension": 0,
        "color": "#ff4444",
    },
    "combinatorics": {
        "file": ROOT / "combinatorics_kernel_v1.json",
        "label": "Combinatorics",
        "dimension": 1,
        "color": "#ff8800",
    },
    "obscure_math": {
        "file": ROOT / "obscure_math_kernel_v1.json",
        "label": "Obscure / Niche Math",
        "dimension": 2,
        "color": "#88cc00",
    },
    "domain": {
        "file": ROOT / "domain_kernel_v1.json",
        "label": "Math Education Domains",
        "dimension": 3,
        "color": "#00cc88",
    },
    "theorem": {
        "file": ROOT / "theorem_kernel_v1.json",
        "label": "Theorem QA",
        "dimension": 4,
        "color": "#0088ff",
    },
    "webmath": {
        "file": ROOT / "webmath_kernel_v1.json",
        "label": "Web Math Patterns",
        "dimension": 5,
        "color": "#8844ff",
    },
}

DIMENSION_DESCRIPTIONS = {
    0: "Diophantine core — Baker bounds, finiteness, tight constraints",
    1: "Combinatorics — additive, extremal, algebraic methods",
    2: "Obscure math — niche subfields, emerging connections",
    3: "Math education — structured domain taxonomy",
    4: "Theorem QA — formal theorem statements",
    5: "Web math — noisy, high-coverage web patterns",
}


def arxiv_check(pid: str) -> bool:
    try:
        r = subprocess.run([
            "ssh", NEON_HOST,
            f"podman exec -i {CONTAINER} psql -U postgres -d {DB} -t -A"
        ], input=f"SELECT 1 FROM arxiv_papers WHERE paper_id = '{pid}'",
           capture_output=True, text=True, timeout=10)
        return r.stdout.strip() == "1"
    except Exception:
        return False


def load_kernel(name: str, source: dict) -> dict | None:
    path = source["file"]
    if not path.exists():
        return None
    return {
        "name": name,
        "label": source["label"],
        "dimension": source["dimension"],
        "color": source["color"],
        "data": json.loads(path.read_text()),
    }


def main():
    print("=" * 60, file=sys.stderr)
    print("Domain Manifold Graph Builder", file=sys.stderr)
    print("=" * 60, file=sys.stderr)

    kernels = {}
    for name, src in KERNEL_SOURCES.items():
        k = load_kernel(name, src)
        if k:
            kernels[name] = k
            print(f"  Loaded {name:20s} dim={k['dimension']}", file=sys.stderr)

    node_map = {}
    edges = []
    added_pairs = set()

    def add_edge(src, tgt, data):
        pair = tuple(sorted([src, tgt]))
        if pair not in added_pairs:
            added_pairs.add(pair)
            edges.append({"source": src, "target": tgt, **data})

    # ── Build nodes from all kernels ──
    for name, kernel in kernels.items():
        d = kernel["data"]
        dim = kernel["dimension"]
        color = kernel["color"]

        # Diophantine: equation_types
        if "equation_types" in d:
            for et in d["equation_types"]:
                nid = f"{name}:{et['type']}"
                node_map[nid] = {
                    "id": nid, "label": et.get("label", et["type"]),
                    "kernel": name, "dimension": dim, "color": color,
                    "type": "equation_type", "papers": et.get("papers", []),
                }

        # Combinatorics: subfields
        if "subfields" in d:
            for sf in d["subfields"]:
                nid = f"{name}:{sf['name']}"
                node_map[nid] = {
                    "id": nid, "label": sf.get("label", sf["name"]),
                    "kernel": name, "dimension": dim, "color": color,
                    "type": "subfield", "papers": sf.get("papers", []),
                }

        # Obscure: domains with name/label
        if "domains" in d and isinstance(d["domains"], list) and d["domains"] and "name" in d["domains"][0]:
            for dom in d["domains"]:
                nid = f"{name}:{dom['name']}"
                node_map[nid] = {
                    "id": nid, "label": dom.get("label", dom["name"]),
                    "kernel": name, "dimension": dim, "color": color,
                    "type": "niche", "papers": dom.get("paper_ids", []),
                }

        # Domain kernel: domains with path/keywords
        if "domains" in d and isinstance(d["domains"], list) and d["domains"] and "path" in d["domains"][0]:
            for dom in d["domains"]:
                nid = f"{name}:{dom['path'][:40]}"
                node_map[nid] = {
                    "id": nid, "label": dom["path"][:80],
                    "kernel": name, "dimension": dim, "color": color,
                    "type": "domain_path", "count": dom.get("count", 0),
                    "leaf": dom.get("leaf", ""),
                }

        # Theorem: kinds
        if "kinds" in d:
            for k in d["kinds"]:
                nid = f"{name}:{k['kind']}"
                node_map[nid] = {
                    "id": nid, "label": f"TheoremQA: {k['kind']}",
                    "kernel": name, "dimension": dim, "color": color,
                    "type": "theorem_kind", "count": k.get("count", 0),
                }

        # Webmath: top patterns
        if "patterns" in d and isinstance(d["patterns"], list):
            for p in d["patterns"][:10]:
                pat = p.get("pattern", "")[:30]
                if not pat.strip():
                    continue
                nid = f"{name}:{pat[:20]}"
                node_map[nid] = {
                    "id": nid, "label": f"web: {pat}",
                    "kernel": name, "dimension": dim, "color": color,
                    "type": "web_pattern", "count": p.get("count", 0),
                }

    print(f"  Nodes: {len(node_map)}", file=sys.stderr)

    # ── Edge type 1: Shared arxiv papers ──
    print(f"  Checking shared papers across kernels...", file=sys.stderr)
    paper_node = defaultdict(list)
    for nid, node in node_map.items():
        for pid in node.get("papers", []):
            paper_node[pid].append(nid)

    for pid, nids in paper_node.items():
        for i in range(len(nids)):
            for j in range(i + 1, len(nids)):
                n1, n2 = nids[i], nids[j]
                if node_map[n1]["kernel"] != node_map[n2]["kernel"]:
                    exists = arxiv_check(pid)
                    if exists:
                        add_edge(n1, n2, {
                            "type": "shared_paper",
                            "paper_id": pid,
                            "verified": exists,
                        })

    print(f"    Shared paper edges: {sum(1 for e in edges if e['type'] == 'shared_paper')}", file=sys.stderr)

    # ── Edge type 2: RRC manifold route connections ──
    receipt_path = Path("archive/experimental-shim-probes/rrc_equation_classifier_receipt.json")
    if receipt_path.exists():
        receipt = json.loads(receipt_path.read_text())
        route_papers = defaultdict(set)
        for e in receipt["compiled_equations"]:
            rec = e["equation_record"]
            route = rec.get("manifold_route", "unclassified")
            pid = rec.get("arxiv_paper_id", "")
            if pid and route not in ("unclassified", "?"):
                route_papers[route].add(pid)

        routes = list(route_papers.keys())
        for i in range(len(routes)):
            for j in range(i + 1, len(routes)):
                shared = route_papers[routes[i]] & route_papers[routes[j]]
                if shared:
                    add_edge(f"manifold:{routes[i]}", f"manifold:{routes[j]}", {
                        "type": "manifold_shared_paper",
                        "shared_count": len(shared),
                        "shared_papers": list(shared)[:5],
                    })

    # ── Edge type 3: Manifold regime connections ──
    if receipt_path.exists():
        receipt = json.loads(receipt_path.read_text())
        regime_papers = defaultdict(set)
        for e in receipt["compiled_equations"]:
            rec = e["equation_record"]
            regime = rec.get("manifold_regime", "diophantine")
            pid = rec.get("arxiv_paper_id", "")
            if pid:
                regime_papers[regime].add(pid)

        regimes = list(regime_papers.keys())
        for i in range(len(regimes)):
            for j in range(i + 1, len(regimes)):
                shared = regime_papers[regimes[i]] & regime_papers[regimes[j]]
                add_edge(f"regime:{regimes[i]}", f"regime:{regimes[j]}", {
                    "type": "regime_shared_paper",
                    "regime_a": regimes[i],
                    "regime_b": regimes[j],
                    "shared_paper_count": len(shared),
                })

    # ── Edge type 4: Shared keywords between kernel types ──
    keyword_sets = {}
    for name, kernel in kernels.items():
        words = set()
        d = kernel["data"]
        if "equation_types" in d:
            for et in d["equation_types"]:
                words.add(et["type"])
                words.add(et.get("label", "").lower())
        if "subfields" in d:
            for sf in d["subfields"]:
                words.add(sf["name"])
                if "methods" in sf:
                    for m in sf["methods"]:
                        words.add(m.lower())
        if "domains" in d and isinstance(d["domains"], list) and d["domains"] and "name" in d["domains"][0]:
            for dom in d["domains"]:
                words.add(dom["name"])
        keyword_sets[name] = words

    k_names = list(keyword_sets.keys())
    for i in range(len(k_names)):
        for j in range(i + 1, len(k_names)):
            shared = keyword_sets[k_names[i]] & keyword_sets[k_names[j]]
            if shared:
                add_edge(f"kernel:{k_names[i]}", f"kernel:{k_names[j]}", {
                    "type": "shared_topic",
                    "shared_terms": list(shared)[:10],
                    "overlap_count": len(shared),
                })

    # ── Edge type 5: Domain adapter bridges (cross-kernel connections) ──
    adapter_path = ROOT / "domain_adapter_kernel_v1.json"
    if adapter_path.exists():
        adapter = json.loads(adapter_path.read_text())
        for b in adapter.get("bridges", []):
            bridge_label = b["bridge"].replace(" ↔ ", "_")
            add_edge(f"adapter:{bridge_label}", f"paper:{b['paper_id']}", {
                "type": "domain_adapter",
                "bridge": b["bridge"],
                "paper_id": b["paper_id"],
                "paper_title": b.get("title", ""),
            })

    # ── Edge type 6: Dimension adjacency (chain connecting dimensions 0→1→2→3→4→5) ──
    for dim in range(5):
        src_nodes = [n for n in node_map.values() if n["dimension"] == dim]
        tgt_nodes = [n for n in node_map.values() if n["dimension"] == dim + 1]
        if src_nodes and tgt_nodes:
            add_edge(f"dim:{dim}", f"dim:{dim + 1}", {
                "type": "dimension_chain",
                "from_dim": dim,
                "to_dim": dim + 1,
                "from_label": DIMENSION_DESCRIPTIONS.get(dim, ""),
                "to_label": DIMENSION_DESCRIPTIONS.get(dim + 1, ""),
            })

    # ── Edge type 7: Intrinsic kernel_internal (content nodes within same kernel) ──
    kernel_groups: dict[str, list[dict]] = defaultdict(list)
    for n in node_map.values():
        kernel_groups[n["kernel"]].append(n)
    for kernel, members in kernel_groups.items():
        if len(members) < 2:
            continue
        for i in range(len(members)):
            for j in range(i + 1, len(members)):
                if i * len(members) + j % 7 == 0:  # ~14% density, deterministic
                    add_edge(members[i]["id"], members[j]["id"], {
                        "type": "kernel_internal",
                        "kernel": kernel,
                        "dimension": members[i].get("dimension"),
                    })

    # ── Edge type 8: kernel_hub (cross-kernel hub connections) ──
    hubs: dict[str, str] = {}
    for kernel, members in kernel_groups.items():
        hubs[kernel] = members[len(members) // 2]["id"]
    kernel_names = list(kernel_groups.keys())
    for i in range(len(kernel_names)):
        for j in range(i + 1, len(kernel_names)):
            add_edge(hubs[kernel_names[i]], hubs[kernel_names[j]], {
                "type": "kernel_hub",
                "from_kernel": kernel_names[i],
                "to_kernel": kernel_names[j],
            })

    # ── Edge type 9: topic_overlap (label token overlap between content nodes) ──
    node_list = list(node_map.values())
    for i in range(len(node_list)):
        ni = node_list[i]
        ni_tokens = set(ni.get("label", "").lower().split())
        for j in range(i + 1, len(node_list)):
            nj = node_list[j]
            overlap = ni_tokens & set(nj.get("label", "").lower().split())
            if len(overlap) >= 2:
                add_edge(ni["id"], nj["id"], {
                    "type": "topic_overlap",
                    "overlap_count": len(overlap),
                    "shared_terms": list(overlap)[:5],
                })

    # ── Edge type 10: shared_paper fallback (unverified) ──
    live_share_count = sum(1 for e in edges if e["type"] == "shared_paper")
    if live_share_count == 0:
        for pid, nids in paper_node.items():
            for i in range(len(nids)):
                for j in range(i + 1, len(nids)):
                    n1, n2 = nids[i], nids[j]
                    if node_map[n1]["kernel"] != node_map[n2]["kernel"]:
                        add_edge(n1, n2, {
                            "type": "shared_paper",
                            "paper_id": pid,
                            "verified": False,
                        })

    # ── Build output ──
    manifold = {
        "schema": "domain_manifold_graph_v1",
        "description": "Expanding manifold graph across all gathered math domains",
        "dimensions": DIMENSION_DESCRIPTIONS,
        "nodes": list(node_map.values()),
        "edges": edges,
        "stats": {
            "total_nodes": len(node_map),
            "total_edges": len(edges),
            "by_type": {t: sum(1 for e in edges if e["type"] == t)
                        for t in set(e["type"] for e in edges)},
        },
    }

    OUT_PATH.parent.mkdir(parents=True, exist_ok=True)
    OUT_PATH.write_text(json.dumps(manifold, indent=2, ensure_ascii=False))

    print(f"\n{'='*60}", file=sys.stderr)
    print(f"Manifold Graph Summary", file=sys.stderr)
    stats = manifold["stats"]
    print(f"  Nodes: {stats['total_nodes']}", file=sys.stderr)
    print(f"  Edges: {stats['total_edges']}", file=sys.stderr)
    print(f"  By type: {stats['by_type']}", file=sys.stderr)

    print(f"\n  Dimensions:", file=sys.stderr)
    for dim, desc in sorted(DIMENSION_DESCRIPTIONS.items()):
        count = sum(1 for n in node_map.values() if n["dimension"] == dim)
        print(f"    Dim {dim}: {desc} ({count} nodes)", file=sys.stderr)

    print(f"\n  Sample edges:", file=sys.stderr)
    for e in edges[:8]:
        print(f"    {e['source']:35s} → {e['target']:35s}  [{e['type']}]", file=sys.stderr)

    print(f"\nSaved to {OUT_PATH}", file=sys.stderr)


if __name__ == "__main__":
    main()