#!/usr/bin/env python3
"""
rrc_root_system_probe.py — Which root system does the RRC equation corpus sit on?

Pipeline:
  1. Fingerprint each equation into a symbol-ROLE vector (15 axes) via the
     math-symbol matrix (math_symbols.CHAR_INFO) on normalize_math(text).
  2. Build the role COUPLING matrix C = corr(role_i, role_j) across the corpus.
  3. Spectrum of C → effective rank (participation ratio) = the data's intrinsic
     dimension.
  4. Minimum spanning tree of roles under distance 1-|corr|. Dynkin diagrams are
     TREES, so the coupling skeleton's tree-topology is directly comparable to
     the A / D / E_n Dynkin diagrams:
        max-degree ≤ 2            → A_n  (path)
        one deg-3 node, arms 1,1,k → D_n  (fork)
        one deg-3 node, arms 1,2,2 → E6
        one deg-3 node, arms 1,2,3 → E7
        one deg-3 node, arms 1,2,4 → E8
  5. Emits a receipt; the per-equation role-vectors + coupling tree are the probe
     we can then apply to NEW math (anomaly = lands in a coupling-tree "hole").

HONEST SCOPE: this matches the coupling *skeleton* to Dynkin *trees* (a real
graph comparison) and reports the effective rank. It does NOT claim C is a Cartan
matrix; that stronger claim would need the ±1/±2 integer inner-product spectrum.

Usage:  python3 4-Infrastructure/shim/rrc_root_system_probe.py
"""
from __future__ import annotations

import json
import os
import sys
import time
from collections import Counter
from pathlib import Path

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from math_symbols import CHAR_INFO, normalize_math  # noqa: E402

ROOT = Path("/home/allaun/Research Stack")
RECEIPT_IN = ROOT / "archive/experimental-shim-probes/rrc_equation_classifier_receipt.json"
RECEIPT_OUT = ROOT / "shared-data/data/rrc_root_system_probe_receipt.json"

# ADE Dynkin diagrams as (branch arm-length signatures). Arms measured in EDGES
# from the unique degree-3 node; A_n has no branch node.
DYNKIN = {
    (1, 1): "D",        # fork: two length-1 arms + a tail of length k  → D_{k+2}
    (1, 2, 2): "E6",
    (1, 2, 3): "E7",
    (1, 2, 4): "E8",
}


def role_vector(text: str, roles: list[str]) -> np.ndarray:
    """Count symbol roles in an equation after LaTeX/Unicode normalization."""
    norm = normalize_math(text)
    idx = {r: i for i, r in enumerate(roles)}
    v = np.zeros(len(roles))
    for ch in norm:
        info = CHAR_INFO.get(ch)
        if info and info["role"] in idx:
            v[idx[info["role"]]] += 1.0
    return v


def load_equations() -> list[str]:
    d = json.loads(RECEIPT_IN.read_text())
    out = []
    for e in d.get("compiled_equations", []):
        r = e.get("equation_record", {})
        txt = (r.get("name", "") + " " + r.get("equation", "")).strip()
        if txt:
            out.append(txt)
    return out


def mst_prim(dist: np.ndarray) -> list[tuple[int, int, float]]:
    """Prim's MST on a dense distance matrix → list of (i, j, dist) edges."""
    n = dist.shape[0]
    in_tree = [False] * n
    in_tree[0] = True
    edges = []
    for _ in range(n - 1):
        best = (None, None, np.inf)
        for i in range(n):
            if not in_tree[i]:
                continue
            for j in range(n):
                if in_tree[j]:
                    continue
                if dist[i, j] < best[2]:
                    best = (i, j, dist[i, j])
        i, j, w = best
        if j is None:
            break
        in_tree[j] = True
        edges.append((i, j, float(w)))
    return edges


def classify_tree(n: int, adj: dict[int, list[int]]) -> tuple[str, dict]:
    """Classify a tree's topology against the A/D/E Dynkin families."""
    deg = {v: len(adj[v]) for v in adj}
    branch = [v for v in deg if deg[v] >= 3]
    info = {"n_nodes": n, "max_degree": max(deg.values()) if deg else 0,
            "n_branch_nodes": len(branch)}
    if info["max_degree"] <= 2:
        info["arms"] = []
        return f"A_{n}", info
    if len(branch) != 1 or info["max_degree"] != 3:
        info["arms"] = []
        return "irregular (not simply-laced ADE tree)", info

    b = branch[0]
    # measure each arm length (in edges) from the branch node out to a leaf
    arms = []
    for start in adj[b]:
        length, prev, cur = 1, b, start
        while True:
            nxts = [x for x in adj[cur] if x != prev]
            if len(nxts) != 1:   # leaf (0) or another branch (>1)
                break
            prev, cur = cur, nxts[0]
            length += 1
        arms.append(length)
    arms.sort()
    info["arms"] = arms
    key2 = tuple(arms[:2])
    if tuple(arms) in DYNKIN:
        return f"{DYNKIN[tuple(arms)]}", info
    if key2 == (1, 1):
        return f"D_{n}", info
    return f"branched (arms={arms}; nearest E-series by long arm)", info


def main() -> None:
    eqs = load_equations()
    # roles present in the matrix, ordered by global frequency for readability
    all_roles = sorted({i["role"] for i in CHAR_INFO.values()})

    M = np.array([role_vector(e, all_roles) for e in eqs])  # (N, R)
    present = M.sum(axis=0) > 0
    roles = [r for r, p in zip(all_roles, present) if p]
    M = M[:, present]
    totals = M.sum(axis=0)

    # drop zero-variance columns (corr undefined)
    var = M.var(axis=0)
    keep = var > 1e-9
    roles = [r for r, k in zip(roles, keep) if k]
    M = M[:, keep]
    R = M.shape[1]

    C = np.corrcoef(M, rowvar=False)
    eig = np.sort(np.linalg.eigvalsh(C))[::-1]
    pos = eig[eig > 1e-9]
    participation = (pos.sum() ** 2) / (np.square(pos).sum())  # effective rank

    dist = 1.0 - np.abs(C)
    np.fill_diagonal(dist, 0.0)
    edges = mst_prim(dist)
    adj: dict[int, list[int]] = {i: [] for i in range(R)}
    for i, j, _ in edges:
        adj[i].append(j)
        adj[j].append(i)
    dynkin, tinfo = classify_tree(R, adj)

    # ---- report ----
    print("=" * 66)
    print(f"RRC ROOT-SYSTEM PROBE  —  {len(eqs)} equations, {R} active role axes")
    print("=" * 66)
    print("\nRole totals across corpus:")
    for r, t in sorted(zip(roles, totals[keep] if keep.shape[0] == totals.shape[0] else totals),
                       key=lambda x: -x[1]):
        print(f"  {r:16s} {int(t)}")
    print(f"\nCoupling-matrix eigenvalues (desc): "
          + ", ".join(f"{e:.2f}" for e in eig))
    print(f"Effective rank (participation ratio): {participation:.2f}  "
          f"→ compare E6=6, E7=7, E8=8")
    print("\nRole coupling MST (Dynkin skeleton):")
    for i, j, w in edges:
        print(f"  {roles[i]:16s} —{1-w:+.2f}— {roles[j]}")
    print(f"\nTree topology: nodes={tinfo['n_nodes']}, max_degree={tinfo['max_degree']}, "
          f"branch_nodes={tinfo['n_branch_nodes']}, arms={tinfo.get('arms')}")
    print(f"\n>>> NEAREST DYNKIN TYPE:  {dynkin}")
    print(f">>> effective dim {participation:.2f} vs rank({R}) skeleton {dynkin}")

    receipt = {
        "schema": "rrc_root_system_probe_v1",
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
        "n_equations": len(eqs),
        "roles": roles,
        "role_totals": {r: int(t) for r, t in zip(roles, totals[keep])},
        "coupling_eigenvalues": [float(x) for x in eig],
        "effective_rank": float(participation),
        "mst_edges": [[roles[i], roles[j], float(1 - w)] for i, j, w in edges],
        "tree_topology": tinfo,
        "nearest_dynkin": dynkin,
        "caveat": "coupling-skeleton vs Dynkin-tree topology; not a Cartan-matrix claim",
    }
    RECEIPT_OUT.write_text(json.dumps(receipt, indent=2, ensure_ascii=False))
    print(f"\nReceipt: {RECEIPT_OUT}")


if __name__ == "__main__":
    main()