epub2md/epub2md/mathml.py

"""Conservative MathML -> TeX conversion and TeX detection helpers.

The converter only handles constructs it can translate unambiguously. Anything
else raises :class:`Unsupported`; callers must then keep the original MathML
instead of emitting possibly wrong TeX.
"""

from __future__ import annotations

import html
import re
from typing import Any

from .utils import is_element, local, norm_ws, text_content


class Unsupported(Exception):
    pass


GREEK = {
    "α": r"\alpha", "β": r"\beta", "γ": r"\gamma", "δ": r"\delta", "ε": r"\varepsilon", "ϵ": r"\epsilon",
    "ζ": r"\zeta", "η": r"\eta", "θ": r"\theta", "ϑ": r"\vartheta", "ι": r"\iota", "κ": r"\kappa",
    "λ": r"\lambda", "μ": r"\mu", "ν": r"\nu", "ξ": r"\xi", "π": r"\pi", "ϖ": r"\varpi", "ρ": r"\rho",
    "ϱ": r"\varrho", "σ": r"\sigma", "ς": r"\varsigma", "τ": r"\tau", "υ": r"\upsilon", "φ": r"\varphi",
    "ϕ": r"\phi", "χ": r"\chi", "ψ": r"\psi", "ω": r"\omega", "Γ": r"\Gamma", "Δ": r"\Delta",
    "Θ": r"\Theta", "Λ": r"\Lambda", "Ξ": r"\Xi", "Π": r"\Pi", "Σ": r"\Sigma", "Υ": r"\Upsilon",
    "Φ": r"\Phi", "Ψ": r"\Psi", "Ω": r"\Omega",
}
OPERATORS = {
    "∑": r"\sum", "∏": r"\prod", "∫": r"\int", "∬": r"\iint", "∮": r"\oint", "⋃": r"\bigcup", "⋂": r"\bigcap",
    "≤": r"\leq", "≥": r"\geq", "≠": r"\neq", "≈": r"\approx", "≡": r"\equiv", "∼": r"\sim", "≅": r"\cong",
    "∝": r"\propto", "×": r"\times", "⋅": r"\cdot", "·": r"\cdot", "∗": r"\ast", "÷": r"\div", "±": r"\pm",
    "∓": r"\mp", "∘": r"\circ", "⊗": r"\otimes", "⊕": r"\oplus", "→": r"\to", "←": r"\leftarrow",
    "↔": r"\leftrightarrow", "⇒": r"\Rightarrow", "⇐": r"\Leftarrow", "⇔": r"\Leftrightarrow",
    "↦": r"\mapsto", "∈": r"\in", "∉": r"\notin", "∋": r"\ni", "⊂": r"\subset", "⊆": r"\subseteq",
    "⊃": r"\supset", "⊇": r"\supseteq", "∪": r"\cup", "∩": r"\cap", "∖": r"\setminus", "∅": r"\emptyset",
    "∀": r"\forall", "∃": r"\exists", "¬": r"\neg", "∧": r"\wedge", "∨": r"\vee", "∞": r"\infty",
    "∂": r"\partial", "∇": r"\nabla", "√": r"\surd", "…": r"\ldots", "⋯": r"\cdots", "⋮": r"\vdots",
    "⋱": r"\ddots", "′": "'", "″": "''", "−": "-", "∣": r"\mid", "∥": r"\parallel", "⊥": r"\perp",
    "⟨": r"\langle", "⟩": r"\rangle", "〈": r"\langle", "〉": r"\rangle", "⌊": r"\lfloor", "⌋": r"\rfloor",
    "⌈": r"\lceil", "⌉": r"\rceil", "ℓ": r"\ell", "ℏ": r"\hbar", "ℵ": r"\aleph", "°": r"^\circ",
    "{": r"\{", "}": r"\}", "|": "|", "‖": r"\|", "\u2061": "", "\u2062": "", "\u2063": ",", "\u2064": "+",
}
FUNCTIONS = {"sin", "cos", "tan", "cot", "sec", "csc", "arcsin", "arccos", "arctan", "sinh", "cosh", "tanh",
             "log", "ln", "lg", "exp", "lim", "liminf", "limsup", "max", "min", "sup", "inf", "det", "dim",
             "ker", "deg", "gcd", "arg", "Pr", "hom"}
LARGE_OPS = {r"\sum", r"\prod", r"\int", r"\iint", r"\oint", r"\bigcup", r"\bigcap", r"\lim", r"\max", r"\min",
             r"\sup", r"\inf", r"\limsup", r"\liminf"}
ACCENTS = {"^": r"\hat", "ˆ": r"\hat", "¯": r"\bar", "‾": r"\overline", "_": r"\overline", "→": r"\vec",
           "⃗": r"\vec", "~": r"\tilde", "˜": r"\tilde", "˙": r"\dot", "¨": r"\ddot", "⏞": r"\overbrace",
           "︷": r"\overbrace"}
UNDER_ACCENTS = {"_": r"\underline", "‾": r"\underline", "¯": r"\underline", "⏟": r"\underbrace", "︸": r"\underbrace"}
VARIANTS = {"bold": r"\mathbf", "double-struck": r"\mathbb", "script": r"\mathcal", "fraktur": r"\mathfrak",
            "normal": r"\mathrm", "sans-serif": r"\mathsf", "monospace": r"\mathtt", "bold-italic": r"\boldsymbol"}
TEX_SPECIAL = {"#": r"\#", "$": r"\$", "%": r"\%", "&": r"\&", "_": r"\_", "~": r"\sim", "^": r"\hat{}", "\\": r"\backslash"}
OPEN_FENCES = {"(": "(", "[": "[", "{": r"\{", "|": "|", "‖": r"\|", "⟨": r"\langle", "⌊": r"\lfloor", "⌈": r"\lceil"}
CLOSE_FENCES = {")": ")", "]": "]", "}": r"\}", "|": "|", "‖": r"\|", "⟩": r"\rangle", "⌋": r"\rfloor", "⌉": r"\rceil"}


def _kids(el: Any) -> list[Any]:
    return [c for c in el if is_element(c)]


def _txt(el: Any) -> str:
    return norm_ws(text_content(el))


def _escape_symbol(ch: str) -> str:
    if ch in GREEK:
        return GREEK[ch]
    if ch in OPERATORS:
        return OPERATORS[ch]
    if ch in TEX_SPECIAL:
        return TEX_SPECIAL[ch]
    return ch


def _symbols(text: str) -> str:
    out = []
    for ch in text:
        sym = _escape_symbol(ch)
        if sym.startswith("\\") and sym[-1].isalpha():
            sym += " "
        out.append(sym)
    return "".join(out).strip()


def _group(tex: str) -> str:
    tex = tex.strip()
    if len(tex) == 1 or re.fullmatch(r"\\[A-Za-z]+", tex) or re.fullmatch(r"\d", tex):
        return tex
    return "{" + tex + "}"


def _text_escape(text: str) -> str:
    return "".join({"\\": r"\textbackslash{}", "{": r"\{", "}": r"\}", "$": r"\$", "&": r"\&",
                    "#": r"\#", "%": r"\%", "_": r"\_", "^": r"\^{}", "~": r"\~{}"}.get(c, c) for c in text)


def tex_annotation(math: Any) -> str | None:
    for el in math.iter():
        if is_element(el) and local(el.tag) == "annotation":
            enc = (el.get("encoding") or "").lower()
            if enc in ("application/x-tex", "tex", "latex", "application/x-latex", "text/x-tex"):
                value = (el.text or "").strip()
                if value:
                    return value
    alttext = math.get("alttext")
    if alttext and looks_like_tex(alttext):
        return strip_tex_delimiters(alttext)
    return None


def mathml_to_tex(math: Any) -> str:
    """Convert a <math> element. Raises Unsupported for constructs we do not handle."""
    annotated = tex_annotation(math)
    if annotated:
        return annotated
    tex = _row(_kids(math))
    tex = re.sub(r"[ ]{2,}", " ", tex).strip()
    if not tex:
        raise Unsupported("empty")
    return tex


def _row(children: list[Any]) -> str:
    # \left ... \right around tall content
    if len(children) >= 3 and local(children[0].tag) == "mo" and local(children[-1].tag) == "mo":
        o, c = _txt(children[0]), _txt(children[-1])
        if o in OPEN_FENCES and c in CLOSE_FENCES and any(_is_tall(k) for k in children[1:-1]):
            inner = " ".join(_node(k) for k in children[1:-1])
            return rf"\left{OPEN_FENCES[o]} {inner} \right{CLOSE_FENCES[c]}"
    return " ".join(_node(k) for k in children)


def _is_tall(el: Any) -> bool:
    return any(is_element(e) and local(e.tag) in ("mfrac", "mtable", "munderover", "msqrt", "mroot") for e in el.iter())


def _node(el: Any) -> str:
    name = local(el.tag)
    handler = _HANDLERS.get(name)
    if handler is None:
        raise Unsupported(name)
    return handler(el)


def _mi(el: Any) -> str:
    text = _txt(el)
    if not text:
        return ""
    variant = el.get("mathvariant", "")
    if len(text) > 1:
        if text in FUNCTIONS:
            return "\\" + text + " "
        base = _text_escape(text)
        cmd = VARIANTS.get(variant, "\\mathrm") if variant not in ("", "normal", "italic") else "\\mathrm"
        return f"{cmd}{{{base}}}"
    sym = _symbols(text)
    if variant and variant != "italic":
        if variant not in VARIANTS:
            raise Unsupported(f"mathvariant={variant}")
        return f"{VARIANTS[variant]}{{{sym}}}"
    return sym


def _mn(el: Any) -> str:
    return _symbols(_txt(el))


def _mo(el: Any) -> str:
    text = _txt(el)
    if text in FUNCTIONS:
        return "\\" + text + " "
    return _symbols(text)


def _mtext(el: Any) -> str:
    text = text_content(el)
    if not text.strip():
        return r"\ " if text else ""
    return rf"\text{{{_text_escape(norm_ws(text))}}}"


def _mspace(el: Any) -> str:
    return r"\,"


def _mrow(el: Any) -> str:
    return _row(_kids(el))


def _first_child(el: Any) -> str:
    kids = _kids(el)
    return _node(kids[0]) if kids else ""


def _mphantom(el: Any) -> str:
    return rf"\phantom{{{_row(_kids(el))}}}"


def _mfrac(el: Any) -> str:
    kids = _kids(el)
    if len(kids) != 2:
        raise Unsupported("mfrac arity")
    lt = (el.get("linethickness") or "").strip()
    if lt in ("0", "0px", "0em", "0pt"):
        return rf"\genfrac{{}}{{}}{{0pt}}{{}}{{{_node(kids[0])}}}{{{_node(kids[1])}}}"
    if el.get("bevelled") == "true":
        return f"{_group(_node(kids[0]))}/{_group(_node(kids[1]))}"
    return rf"\frac{{{_node(kids[0])}}}{{{_node(kids[1])}}}"


def _msqrt(el: Any) -> str:
    return rf"\sqrt{{{_row(_kids(el))}}}"


def _mroot(el: Any) -> str:
    kids = _kids(el)
    if len(kids) != 2:
        raise Unsupported("mroot arity")
    return rf"\sqrt[{_node(kids[1])}]{{{_node(kids[0])}}}"


def _base(el: Any) -> str:
    tex = _node(el)
    if local(el.tag) in ("mi", "mn", "mo") and len(_txt(el)) == 1:
        return tex.strip()
    if re.fullmatch(r"\\[A-Za-z]+\s*", tex):
        return tex.strip()
    return "{" + tex + "}"


def _msup(el: Any) -> str:
    kids = _kids(el)
    if len(kids) != 2:
        raise Unsupported("msup arity")
    return f"{_base(kids[0])}^{{{_node(kids[1])}}}"


def _msub(el: Any) -> str:
    kids = _kids(el)
    if len(kids) != 2:
        raise Unsupported("msub arity")
    return f"{_base(kids[0])}_{{{_node(kids[1])}}}"


def _msubsup(el: Any) -> str:
    kids = _kids(el)
    if len(kids) != 3:
        raise Unsupported("msubsup arity")
    return f"{_base(kids[0])}_{{{_node(kids[1])}}}^{{{_node(kids[2])}}}"


def _mover(el: Any) -> str:
    kids = _kids(el)
    if len(kids) != 2:
        raise Unsupported("mover arity")
    base, over = kids
    over_txt = _txt(over)
    if local(over.tag) == "mo" and over_txt in ACCENTS:
        return rf"{ACCENTS[over_txt]}{{{_node(base)}}}"
    base_tex = _node(base).strip()
    if base_tex in LARGE_OPS:
        return f"{base_tex}^{{{_node(over)}}}"
    return rf"\overset{{{_node(over)}}}{{{base_tex}}}"


def _munder(el: Any) -> str:
    kids = _kids(el)
    if len(kids) != 2:
        raise Unsupported("munder arity")
    base, under = kids
    under_txt = _txt(under)
    if local(under.tag) == "mo" and under_txt in UNDER_ACCENTS:
        return rf"{UNDER_ACCENTS[under_txt]}{{{_node(base)}}}"
    base_tex = _node(base).strip()
    if base_tex in LARGE_OPS:
        return f"{base_tex}_{{{_node(under)}}}"
    return rf"\underset{{{_node(under)}}}{{{base_tex}}}"


def _munderover(el: Any) -> str:
    kids = _kids(el)
    if len(kids) != 3:
        raise Unsupported("munderover arity")
    base_tex = _node(kids[0]).strip()
    if base_tex in LARGE_OPS:
        return f"{base_tex}_{{{_node(kids[1])}}}^{{{_node(kids[2])}}}"
    return rf"\overset{{{_node(kids[2])}}}{{\underset{{{_node(kids[1])}}}{{{base_tex}}}}}"


def _mfenced(el: Any) -> str:
    o = el.get("open", "(")
    c = el.get("close", ")")
    seps = (el.get("separators", ",") or "").replace(" ", "")
    kids = _kids(el)
    parts = []
    for i, k in enumerate(kids):
        parts.append(_node(k))
        if i < len(kids) - 1 and seps:
            parts.append(_symbols(seps[min(i, len(seps) - 1)]))
    lo = OPEN_FENCES.get(o, "." if o == "" else None)
    rc = CLOSE_FENCES.get(c, "." if c == "" else None)
    if lo is None or rc is None:
        raise Unsupported(f"mfenced {o}{c}")
    return rf"\left{lo} {' '.join(parts)} \right{rc}"


def _mtable(el: Any) -> str:
    rows = []
    for tr in _kids(el):
        if local(tr.tag) == "mlabeledtr":
            raise Unsupported("mlabeledtr")
        if local(tr.tag) != "mtr":
            raise Unsupported(local(tr.tag))
        cells = []
        for td in _kids(tr):
            if local(td.tag) != "mtd":
                raise Unsupported(local(td.tag))
            cells.append(_row(_kids(td)))
        rows.append(" & ".join(cells))
    return r"\begin{matrix} " + r" \\ ".join(rows) + r" \end{matrix}"


def _semantics(el: Any) -> str:
    kids = [k for k in _kids(el) if local(k.tag) not in ("annotation", "annotation-xml")]
    if not kids:
        raise Unsupported("semantics")
    return _node(kids[0])


def _mstyle(el: Any) -> str:
    return _row(_kids(el))


_HANDLERS = {
    "mi": _mi, "mn": _mn, "mo": _mo, "mtext": _mtext, "ms": _mtext, "mspace": _mspace,
    "mrow": _mrow, "mstyle": _mstyle, "mpadded": _mrow, "merror": _mrow, "mphantom": _mphantom,
    "mfrac": _mfrac, "msqrt": _msqrt, "mroot": _mroot, "msup": _msup, "msub": _msub,
    "msubsup": _msubsup, "mover": _mover, "munder": _munder, "munderover": _munderover,
    "mfenced": _mfenced, "mtable": _mtable, "semantics": _semantics, "maction": _first_child,
    "none": lambda el: "{}",
}


# --------------------------------------------------------------------------- TeX in text/attributes

TEX_HINT = re.compile(r"\\[A-Za-z]+|[\^_]\{|\\[\[\]()]|^\$.*\$$", re.S)


def looks_like_tex(text: str) -> bool:
    t = (text or "").strip()
    if not t:
        return False
    if (t.startswith("$") and t.endswith("$") and len(t) > 2) or t.startswith(("\\(", "\\[")):
        return True
    return bool(re.search(r"\\(frac|sqrt|sum|int|alpha|beta|mathbb|mathrm|begin|left|right|cdot|in|leq|geq)\b", t))


def strip_tex_delimiters(text: str) -> str:
    t = html.unescape(html.unescape((text or "").strip()))  # some EPUBs double-escape '&' in alt texts
    for o, c in (("$$", "$$"), ("\\[", "\\]"), ("\\(", "\\)"), ("$", "$")):
        if t.startswith(o) and t.endswith(c) and len(t) >= len(o) + len(c):
            t = t[len(o) : len(t) - len(c)]
            break
    return portable_tex(t.strip())


UMLAUTS = {"a": "ä", "o": "ö", "u": "ü", "A": "Ä", "O": "Ö", "U": "Ü"}


def portable_tex(tex: str) -> str:
    """Publisher TeX → TeX that Pandoc's formula parser (texmath, MathJax, KaTeX) understands, still valid LaTeX:
    text accents \\"u → ü, \\O (empty set in Springer books) → \\emptyset, no @{…} column separators in
    array specifications, no \\arraycolsep assignments, no empty \\displaystyle before a cell end. Only changes
    spacing, never content."""
    tex = re.sub(r'\\"(?:\{([aouAOU])\}|([aouAOU]))', lambda m: UMLAUTS[m.group(1) or m.group(2)], tex)
    tex = re.sub(r"\\(?:text|mbox|mathrm)\s*\{\s*\\O\s*\}", r"\\emptyset", tex)
    tex = re.sub(r"\\O(?![A-Za-z])", r"\\emptyset", tex)
    tex = re.sub(r"\\arraycolsep\s*-?\d*\.?\d+\s*(?:pt|em|ex|mm|cm)", "", tex)
    tex = re.sub(r"\\displaystyle\s*(?=\\\\|&|\\end\b)", "", tex)  # no-op before a cell end; texmath fails
    out: list[str] = []
    pos = 0
    for m in re.finditer(r"\\begin\{array\}\{", tex):
        if m.start() < pos:
            continue
        end = _closing_brace(tex, m.end() - 1)
        if end < 0:
            break
        spec = tex[m.end():end]
        i = spec.find("@{")
        while i >= 0:
            close = _closing_brace(spec, i + 1)
            if close < 0:
                break
            spec = spec[:i] + spec[close + 1:]
            i = spec.find("@{")
        out.append(tex[pos:m.end()] + spec)
        pos = end
    return "".join(out) + tex[pos:]


def _closing_brace(text: str, open_pos: int) -> int:
    """Index of the '}' matching the '{' at ``open_pos`` (-1 if unbalanced)."""
    depth = 0
    for i in range(open_pos, len(text)):
        if text[i] == "{":
            depth += 1
        elif text[i] == "}":
            depth -= 1
            if depth == 0:
                return i
    return -1


def is_display_tex(text: str) -> bool:
    t = (text or "").strip()
    return t.startswith(("$$", "\\[")) or "\\displaystyle" in t[:40]