#!/usr/bin/env python3
"""
math_symbols.py — Math-symbol database loader + LaTeX/Unicode normalizer.

Backs the RRC geometry / tensor-notation kernel. Provides:

  * MATH_SYMBOLS   — full symbol table from shared-data/data/math_symbols_v1.json
                     (built by build_math_symbols_db.py: unicode-math-table.tex +
                     unicodedata, ~2950 symbols).
  * LATEX_TO_CHAR  — command → Unicode char. A curated map of the STANDARD LaTeX
                     macros (\\Gamma, \\rho, \\nabla, …) takes precedence over the
                     unicode-math table (which binds \\Gamma to a math-italic
                     codepoint, not the plain Greek letter), then the full DB
                     fills in the long tail (\\boxtimes, \\curlyvee, …).
  * CHAR_INFO      — char → {role, name, block, category} for feature extraction.
  * normalize_math(text) — canonicalize LaTeX/Unicode so the notation signatures
                     match regardless of encoding:  \\Gamma ≡ Γ, \\rho\\sigma ≡ ρσ,
                     and Penrose abstract-index  R^{a}{}_{bcd}  collapses cleanly.

Degrades gracefully: if the JSON DB is absent, the curated map alone still drives
normalization of the common macros.
"""
from __future__ import annotations

import json
import re
from pathlib import Path

_DB_PATHS = [
    Path("/home/allaun/Research Stack/shared-data/data/math_symbols_v1.json"),
    Path(__file__).resolve().parent.parent.parent / "shared-data/data/math_symbols_v1.json",
    Path("shared-data/data/math_symbols_v1.json"),
]

# ── Curated standard-LaTeX macros (authoritative for the common commands) ─────
_GREEK = {
    "alpha": "α", "beta": "β", "gamma": "γ", "delta": "δ", "epsilon": "ε",
    "varepsilon": "ε", "zeta": "ζ", "eta": "η", "theta": "θ", "vartheta": "ϑ",
    "iota": "ι", "kappa": "κ", "lambda": "λ", "mu": "μ", "nu": "ν", "xi": "ξ",
    "omicron": "ο", "pi": "π", "varpi": "ϖ", "rho": "ρ", "varrho": "ϱ",
    "sigma": "σ", "varsigma": "ς", "tau": "τ", "upsilon": "υ", "phi": "φ",
    "varphi": "φ", "chi": "χ", "psi": "ψ", "omega": "ω",
    "Gamma": "Γ", "Delta": "Δ", "Theta": "Θ", "Lambda": "Λ", "Xi": "Ξ",
    "Pi": "Π", "Sigma": "Σ", "Upsilon": "Υ", "Phi": "Φ", "Psi": "Ψ", "Omega": "Ω",
}
_OPS = {
    "nabla": "∇", "partial": "∂", "infty": "∞", "times": "×", "cdot": "⋅",
    "otimes": "⊗", "oplus": "⊕", "odot": "⊙", "wedge": "∧", "vee": "∨",
    "pm": "±", "mp": "∓", "ast": "∗", "star": "⋆", "circ": "∘", "bullet": "∙",
    "to": "→", "rightarrow": "→", "longrightarrow": "→", "mapsto": "↦",
    "leftarrow": "←", "Rightarrow": "⇒", "Leftarrow": "⇐", "leftrightarrow": "↔",
    "leq": "≤", "le": "≤", "geq": "≥", "ge": "≥", "neq": "≠", "ne": "≠",
    "approx": "≈", "equiv": "≡", "cong": "≅", "sim": "∼", "simeq": "≃",
    "propto": "∝", "in": "∈", "notin": "∉", "ni": "∋", "subset": "⊂",
    "subseteq": "⊆", "supset": "⊃", "supseteq": "⊇", "cup": "∪", "cap": "∩",
    "setminus": "∖", "emptyset": "∅", "forall": "∀", "exists": "∃",
    "sum": "∑", "prod": "∏", "coprod": "∐", "int": "∫", "oint": "∮", "iint": "∬",
    "Box": "□", "square": "□", "Diamond": "◇", "dagger": "†", "ddagger": "‡",
    "ell": "ℓ", "hbar": "ℏ", "Re": "ℜ", "Im": "ℑ", "aleph": "ℵ", "wp": "℘",
    "angle": "∠", "perp": "⊥", "parallel": "∥", "nparallel": "∦", "top": "⊤",
    "bot": "⊥", "models": "⊨", "vdash": "⊢", "boxtimes": "⊠", "boxplus": "⊞",
    "rtimes": "⋊", "ltimes": "⋉", "bigwedge": "⋀", "bigvee": "⋁",
    "bigcup": "⋃", "bigcap": "⋂", "bigotimes": "⨂", "bigoplus": "⨁", "bigodot": "⨀",
    "langle": "⟨", "rangle": "⟩", "lVert": "‖", "rVert": "‖", "Vert": "‖",
    "nabla": "∇", "triangle": "△", "sharp": "♯", "flat": "♭", "lor": "∨", "land": "∧",
}
# Formatting / spacing macros that carry no symbol meaning — stripped.
_DROP_WORD = {
    "mathrm", "mathbf", "mathit", "mathsf", "mathtt", "mathcal", "mathbb",
    "mathfrak", "mathscr", "boldsymbol", "bm", "operatorname", "text", "textrm",
    "textbf", "textit", "mathnormal", "left", "right", "big", "Big", "bigg",
    "Bigg", "bigl", "bigr", "Bigl", "Bigr", "displaystyle", "textstyle",
    "scriptstyle", "limits", "nolimits", "quad", "qquad",
}

_CURATED: dict[str, str] = {**_GREEK, **_OPS}


def _load_db() -> dict:
    for p in _DB_PATHS:
        try:
            if p.exists():
                return json.loads(p.read_text(encoding="utf-8"))
        except Exception:
            continue
    return {"symbols": []}


MATH_SYMBOLS: list[dict] = _load_db().get("symbols", [])

# command (no backslash) → char.  Curated wins; DB fills the long tail.
LATEX_TO_CHAR: dict[str, str] = {}
for _s in MATH_SYMBOLS:
    _cmd = (_s.get("latex") or "").lstrip("\\").strip()
    if _cmd and _cmd.isalpha() and _cmd not in LATEX_TO_CHAR:
        LATEX_TO_CHAR[_cmd] = _s["char"]
LATEX_TO_CHAR.update(_CURATED)  # curated standard macros take precedence

CHAR_INFO: dict[str, dict] = {s["char"]: s for s in MATH_SYMBOLS}

_CMD_RE = re.compile(r"\\([A-Za-z]+)")
_SPACE_RE = re.compile(r"\\[,!;:> ]")          # \, \! \; \: thin/neg spaces
_EMPTY_GRP_RE = re.compile(r"\{\s*\}")          # Penrose empty index slots {}
_WS_RE = re.compile(r"[ \t]+")


def _sub_cmd(m: re.Match) -> str:
    word = m.group(1)
    if word in _DROP_WORD:
        return " "
    if word in LATEX_TO_CHAR:
        return LATEX_TO_CHAR[word]
    return m.group(0)  # unknown command: leave untouched


def normalize_math(text: str) -> str:
    """Canonicalize LaTeX + Unicode math so notation signatures match uniformly.

    \\Gamma → Γ, \\rho\\sigma\\mu\\nu → ρσμν, \\nabla_\\mu → ∇_μ, strips
    \\mathrm/\\left/\\, wrappers, and collapses Penrose empty index groups
    R^{a}{}_{bcd} → R^{a}_{bcd}. Idempotent on already-Unicode input.
    """
    if not text:
        return ""
    text = _SPACE_RE.sub(" ", text)
    # iterate to resolve nested wrappers like \mathrm{\Gamma}
    for _ in range(3):
        new = _CMD_RE.sub(_sub_cmd, text)
        if new == text:
            break
        text = new
    text = _EMPTY_GRP_RE.sub("", text)
    return _WS_RE.sub(" ", text).strip()