"""Character vocabulary for the CTC recognizer. Index 0 is the CTC blank."""

import re
import string
import unicodedata

LATIN_EXTRA = "ÀÁÂÄÃÅÇÈÉÊËÌÍÎÏÑÒÓÔÖÕØÙÚÛÜÝßàáâäãåçèéêëìíîïñòóôöõøùúûüýÿ"
SYMBOLS = "€£¥₹¢°±×÷§©®™•–—‘’“”…½¼¾μ"   # last is Greek mu (the micro sign folds into it)
# Inline maths, read as plain Unicode text (x², H₂O, α ≤ β). Appended after the original vocabulary so
# models trained before it keep their class indices.
SUPERSCRIPT = dict(zip("0123456789+-=()ni", "⁰¹²³⁴⁵⁶⁷⁸⁹⁺⁻⁼⁽⁾ⁿⁱ"))
SUBSCRIPT = dict(zip("0123456789+-=()aeoxhklmnpstijruv", "₀₁₂₃₄₅₆₇₈₉₊₋₌₍₎ₐₑₒₓₕₖₗₘₙₚₛₜᵢⱼᵣᵤᵥ"))
GREEK = "αβγδεζηθκλνξπρστφχψωΓΔΘΛΞΠΣΦΨΩ"
MATH = "≤≥≠≈≡∞√∫∂∇∈∉⊂⊆∪∩∀∃→←↔⇒⇔∝∠⊥∥"

CHARS = (string.digits + string.ascii_letters + string.punctuation + " " + LATIN_EXTRA + SYMBOLS
         + "".join(SUPERSCRIPT.values()) + "".join(SUBSCRIPT.values()) + GREEK + MATH)
MATH_CHARS = frozenset("".join(SUPERSCRIPT.values()) + "".join(SUBSCRIPT.values()) + GREEK + MATH + "μ")
BLANK = 0
CHAR_TO_ID = {c: i + 1 for i, c in enumerate(CHARS)}
NUM_CLASSES = len(CHARS) + 1
# The field extractor gives every inline-maths character one shared id (it only needs to know one is there).
MATH_ID = CHAR_TO_ID["⁰"]
EXTRACTOR_CLASSES = MATH_ID + 1

# Characters outside the vocabulary that have an obvious in-vocabulary spelling (Greek capitals that
# look Latin are read as Latin).
_REPLACE = {" ": " ", "\t": " ", " ": " ", " ": " ", "‐": "-", "‑": "-", "−": "-",
            "‚": ",", "„": "\"", "′": "'", "″": "\"", "·": ".", "⋅": ".", "▪": "•", "●": "•",
            "µ": "μ", "ϵ": "ε", "ϕ": "φ", "ϑ": "θ", "ς": "σ", "∆": "Δ", "∑": "Σ", "∏": "Π", "∗": "*", "∕": "/",
            "⁄": "/", "∼": "~", "⩽": "≤", "⩾": "≥", "≦": "≤", "≧": "≥", "ο": "o", "ι": "i", "υ": "u",
            **dict(zip("ΑΒΕΖΗΙΚΜΝΟΡΤΥΧ", "ABEZHIKMNOPTYX"))}
_SCRIPT_SPECIAL = {"TM": "™", "tm": "™"}
_SCRIPT_MARKUP = re.compile(r"\$([_^])\{([^{}$]*)\}\$")


def scripted(chars: str, kind: str) -> str:
    """Spell a raised ("^") or lowered ("_") run: in Unicode when every character has a form, else plainly."""
    table = SUPERSCRIPT if kind == "^" else SUBSCRIPT
    if kind == "^" and chars in _SCRIPT_SPECIAL:
        return _SCRIPT_SPECIAL[chars]
    return "".join(table[c] for c in chars) if chars and all(c in table for c in chars) else chars


def latex_scripts_to_unicode(text: str) -> tuple[str, bool]:
    """'H$_{2}$O' -> 'H₂O' (the markup PDF text extraction uses for sub/superscripts).
    Returns (text, clean); clean is False when '$' markup this could not convert is left over."""
    out, n = _SCRIPT_MARKUP.subn(lambda m: scripted(m.group(2).strip(), m.group(1)), text)
    return out, "$" not in out or out.count("$") == text.count("$") - 2 * n and "$_{" not in out and "$^{" not in out


_SUP_RUN = re.compile("[" + "".join(SUPERSCRIPT.values()) + "]+")
_SUB_RUN = re.compile("[" + "".join(SUBSCRIPT.values()) + "]+")
_SUP_BACK = {v: k for k, v in SUPERSCRIPT.items()}
_SUB_BACK = {v: k for k, v in SUBSCRIPT.items()}


def scripts_to_latex(text: str) -> str:
    """'x² + H₂O' -> 'x$^{2}$ + H$_{2}$O', the usual way Markdown writes inline maths."""
    text = _SUP_RUN.sub(lambda m: "$^{" + "".join(_SUP_BACK[c] for c in m.group()) + "}$", text)
    return _SUB_RUN.sub(lambda m: "$_{" + "".join(_SUB_BACK[c] for c in m.group()) + "}$", text)


def normalize(text: str) -> str:
    """Map text into the vocabulary: fold known variants, strip accents only when needed, drop the rest."""
    out = []
    for ch in text:
        ch = _REPLACE.get(ch, ch)
        if ch in CHAR_TO_ID:
            out.append(ch)
            continue
        base = unicodedata.normalize("NFKD", ch)
        out.extend(c for c in base if c in CHAR_TO_ID)
    return " ".join("".join(out).split())


def encode(text: str) -> list[int]:
    return [CHAR_TO_ID[c] for c in text if c in CHAR_TO_ID]


def ctc_greedy_decode(probs, ids=None):
    """probs: (T, C) numpy array of per-step probabilities. Returns (text, per-char confidences)."""
    best = probs.argmax(-1) if ids is None else ids
    text, confs, prev = [], [], BLANK
    for t, k in enumerate(best):
        k = int(k)
        if k != BLANK and k != prev:
            text.append(CHARS[k - 1])
            confs.append(float(probs[t, k]))
        prev = k
    return "".join(text), confs