"""CTC text recognition postprocessor.

Used for: vgg_english_g2 (plate OCR).

Input: [1, seqLen, numChars] logits
Output: {"kind": "text", "text": "ABC123", "confidence": 0.95, "formatValid": true}

Algorithm: argmax per timestep → collapse consecutive duplicates → remove blank (index 0) → join

Plausibility gate
-----------------
When the recognizer is fed a too-low-resolution / heavily-blurred plate crop
its logit map collapses toward the CTC blank and greedy decode emits a spurious
1–2 character read (e.g. a lone "A") at whatever confidence the single surviving
timestep carries. That noise was being stored verbatim as a "plate" label. The
charset + greedy decode themselves are correct — this only rejects a result too
short to be a real plate (`minTextLength`, non-space chars) or below a confidence
floor (`minTextConfidence`), returning empty text so no garbage label is stored.
Real plates read at 5+ chars / conf ~0.86–1.0, so the conservative defaults never
suppress a genuine read; both are config-overridable for per-camera tuning.

Normalisation before the grammar
--------------------------------
The recognizer's charset carries lowercase letters and a space, and live reads
use them freely ("dn Ro 3o9" is the same plate as "DNRO309"). The grammar and
the confusable table are both defined over UPPERCASE alphanumerics, so a raw
read is upper-cased and stripped of non-alphanumerics BEFORE either runs —
otherwise a perfectly readable plate fails the grammar for its case, and the
0↔O repair cannot even be attempted (lowercase 'o' is not a CONFUSABLES key).
The per-character scores and aligned timesteps are filtered alongside the
glyphs, so the rescorer keeps reading the softmax mass at the right timestep.

The RAW text still feeds the symbol gate below: normalisation must never be
able to launder hallucinated punctuation into a conforming plate. A read that
carries a symbol glyph is rejected before the grammar is ever consulted, so
repair only ever runs on a read that already passed the plausibility gate.

Format-aware rescoring (post-hoc)
---------------------------------
After the plausibility gate, the NORMALISED read is validated against the active
region's plate grammar (`plateRegion`, default 'DE'; 'off' disables). If the
letter/digit sequence already matches, it is accepted as-is (`formatValid: true`).
If it does NOT match, candidate reads are generated by substituting classic OCR
confusables (0↔O, 1↔I, 8↔B, 5↔S, 2↔Z, 6↔G, 4↔A) position-by-position; each
candidate is scored by the mean per-character probability and the highest-scoring
GRAMMAR-CONFORMING candidate replaces the read (`formatValid: true`). If no
candidate matches, the original greedy read is kept verbatim (`formatValid:
false`) — validation NEVER drops a plausible read, it only fixes or annotates it.

Per-character probability approximation
---------------------------------------
The greedy CTC path collapses consecutive-duplicate argmax timesteps; for each
emitted output character we retain the single timestep `t` at which its run
began. A candidate's per-character probability is then approximated by the
softmax at that ONE aligned timestep — for an unchanged position the original
argmax probability, and for a substituted position `probs[t, altIndex]` (the
real softmax mass the confusable already carried at that timestep). This is an
approximation: it does not re-run a full CTC forward alignment for the candidate
string (which could re-collapse differently). For single-character confusable
swaps it is a faithful, cheap proxy — the timestep that produced the original
glyph is exactly where its confusable competes for probability.
"""
import itertools
import re

import numpy as np

# Shortest plausible plate (non-space chars). No real plate is 1–2 chars, so
# this only rejects the degenerate near-all-blank collapse; kept conservative to
# never drop a genuine short/vanity plate.
DEFAULT_MIN_TEXT_LENGTH = 3
# Confidence floor (0.0 = disabled by default — character count is the reliable
# discriminator for the degenerate collapse; the "A"@1.0 case is caught by
# length, not confidence). Available for per-camera tuning via config.
DEFAULT_MIN_TEXT_CONFIDENCE = 0.0
# Product-level default plate grammar region — supplied to this postprocessor by
# the TS step config (`plateRegion`, see step-definitions.ts). When the key is
# ABSENT from `config` the postprocessor falls back to 'off' (no rescoring) so a
# bare / pre-plumbing call is byte-identical to the historical decode behaviour;
# the DE default only takes effect once the config plumbing supplies it.
DEFAULT_PLATE_REGION = "DE"
# Fallback when `plateRegion` is absent from config (backward-compatible: no
# rescoring, no formatValid annotation).
ABSENT_PLATE_REGION = "off"
# Cap on the number of positions that may simultaneously carry a confusable
# alternative. Each substitutable position doubles the candidate count, so the
# worst case is 2**MAX_SUBST_POSITIONS candidates. Plates are ≤8 glyphs so this
# is rarely binding, but the cap guarantees a bounded search (≤256 candidates)
# even on a pathological long read: when more positions than this have a
# confusable, only the LOWEST-confidence positions (the most likely misreads)
# stay substitutable; the rest are locked to their greedy glyph.
MAX_SUBST_POSITIONS = 8

# Classic OCR confusable pairs, bidirectional. Maps a glyph to the alternates it
# is commonly mistaken for/against. Used ONLY to repair a read that fails the
# grammar — never to alter a read that already conforms.
CONFUSABLES: dict[str, tuple[str, ...]] = {
    "0": ("O",),
    "O": ("0",),
    "1": ("I",),
    "I": ("1",),
    "8": ("B",),
    "B": ("8",),
    "5": ("S",),
    "S": ("5",),
    "2": ("Z",),
    "Z": ("2",),
    "6": ("G",),
    "G": ("6",),
    "4": ("A",),
    "A": ("4",),
}

# Region plate grammars, keyed by ISO-3166 alpha-2 code. Each pattern validates
# the RAW letter/digit sequence the OCR emits (no separators — the recognizer
# never outputs the dash/space of the physical plate). Ready for future
# countries ('IT', 'FR', …); 'off' (handled before lookup) disables rescoring.
#
# GERMANY ('DE'): district code (1–3 letters) + 1–2 letters + 1–4 digits +
# optional E (electric) / H (historic) suffix; ≤8 glyphs excluding separators.
# Separators are optional in the pattern for robustness, though the OCR emits
# none. The letter-block split (district vs. series) is inherently ambiguous on
# the raw string, but that does not matter — the regex validates the whole
# sequence, not the split.
PLATE_GRAMMARS: dict[str, str] = {
    "DE": r"^[A-ZÄÖÜ]{1,3}[- ]?[A-Z]{1,2}[- ]?[0-9]{1,4}[EH]?$",
}


def _matches_grammar(text: str, pattern) -> bool:
    return pattern.match(text) is not None


def _normalize_read(
    chars: list[str],
    char_scores: list[float],
    char_timesteps: list[int],
) -> tuple[str, list[float], list[int]]:
    """Upper-case + drop non-alphanumeric glyphs, keeping scores/timesteps aligned.

    The grammar patterns and CONFUSABLES are defined over uppercase alnums; the
    charset the recognizer decodes from is not. Returning the filtered score and
    timestep lists alongside the text is what lets the rescorer keep reading a
    candidate glyph's softmax mass at the timestep that produced its position.

    A glyph whose upper-case form is not a single character (no such glyph in
    this charset, but the fold is locale-dependent in general) is kept as-is
    rather than silently changing the string length.
    """
    out_chars: list[str] = []
    out_scores: list[float] = []
    out_timesteps: list[int] = []
    for ch, score, timestep in zip(chars, char_scores, char_timesteps):
        if not ch.isalnum():
            continue
        upper = ch.upper()
        out_chars.append(upper if len(upper) == 1 else ch)
        out_scores.append(score)
        out_timesteps.append(timestep)
    return "".join(out_chars), out_scores, out_timesteps


def _rescore_for_grammar(
    text: str,
    char_scores: list[float],
    char_timesteps: list[int],
    probs: np.ndarray,
    charset: list[str],
    pattern,
) -> tuple[str, bool]:
    """Return (repaired_or_original_text, format_valid).

    If `text` already matches `pattern`, return it unchanged with True.
    Otherwise enumerate confusable substitutions (bounded), score each candidate
    by mean per-character probability, and return the highest-scoring candidate
    that matches the grammar. If none match, return the original text with False.
    """
    if _matches_grammar(text, pattern):
        return text, True

    chars = list(text)
    n = len(chars)

    # Reverse charset lookup (glyph → first index) so a substituted glyph's
    # softmax mass can be read from the aligned timestep. Built once per call.
    char_to_index: dict[str, int] = {}
    for idx, glyph in enumerate(charset):
        if glyph and glyph not in char_to_index:
            char_to_index[glyph] = idx

    # Which positions carry a usable confusable alternative (both the glyph has a
    # mapping AND the alternate is representable in this charset so its timestep
    # probability can be looked up).
    subst_positions: list[int] = []
    for i, ch in enumerate(chars):
        alts = CONFUSABLES.get(ch)
        if not alts:
            continue
        if any(alt in char_to_index for alt in alts):
            subst_positions.append(i)

    if not subst_positions:
        return text, False

    # Bound the search: keep only the MAX_SUBST_POSITIONS lowest-confidence
    # positions substitutable (most likely misreads); lock the rest.
    if len(subst_positions) > MAX_SUBST_POSITIONS:
        subst_positions.sort(key=lambda i: char_scores[i])
        subst_positions = sorted(subst_positions[:MAX_SUBST_POSITIONS])

    subst_set = set(subst_positions)

    # Per-position option list: original glyph, plus any representable confusables
    # (only for substitutable positions). itertools.product enumerates the grid.
    options: list[list[str]] = []
    for i, ch in enumerate(chars):
        if i in subst_set:
            alts = [a for a in CONFUSABLES.get(ch, ()) if a in char_to_index]
            options.append([ch, *alts])
        else:
            options.append([ch])

    def _prob_at(i: int, glyph: str) -> float:
        if glyph == chars[i]:
            return char_scores[i]
        idx = char_to_index.get(glyph)
        if idx is None:
            return 0.0
        return float(probs[char_timesteps[i], idx])

    best_text: str | None = None
    best_score = -1.0
    for combo in itertools.product(*options):
        candidate = "".join(combo)
        if candidate == text:
            continue  # already known non-matching (checked above)
        if not _matches_grammar(candidate, pattern):
            continue
        score = sum(_prob_at(i, combo[i]) for i in range(n)) / n if n else 0.0
        if score > best_score:
            best_score = score
            best_text = candidate

    if best_text is not None:
        return best_text, True
    return text, False


# Glyphs the recognizer emits for a plate's physical registration seal.
SEAL_GLYPHS = frozenset(".:;,*-'`\"·")


def _strip_seal_glyph(text: str, region: str) -> "str | None":
    """Return the normalised plate when `text` is a conforming read carrying ONE
    interior seal glyph; ``None`` otherwise (caller applies the normal gate).

    Narrow by construction: exactly one seal-like glyph, never at either end,
    nothing else non-alphanumeric left over, and the normalised remainder must
    match the region grammar. A read that does not become a legal plate is not
    rescued.
    """
    pattern_src = PLATE_GRAMMARS.get(region)
    if pattern_src is None:
        return None
    seals = [i for i, c in enumerate(text) if c in SEAL_GLYPHS]
    if len(seals) != 1:
        return None
    pos = seals[0]
    if pos == 0 or pos == len(text) - 1:
        return None
    stripped = text[:pos] + text[pos + 1 :]
    if any(not (c.isalnum() or c.isspace()) for c in stripped):
        return None
    normalised = "".join(c for c in stripped.upper() if c.isalnum())
    if re.match(pattern_src, normalised) is None:
        return None
    return normalised


def postprocess_ctc(
    predictions: dict,
    config: dict,
    orig_w: int,
    orig_h: int,
    scale: float,
    pad: tuple[int, int],
) -> dict:
    """CTC greedy decode."""
    charset = config.get("charset", [])

    # Get first output tensor
    raw = np.array(list(predictions.values())[0]).astype(np.float32)

    # Handle shapes: [1, T, C] or [T, C]
    if raw.ndim == 3:
        raw = raw[0]  # remove batch
    # raw is now [T, C] — logits (not probabilities)
    seq_len, num_chars = raw.shape

    # Apply softmax to get probabilities
    raw_max = np.max(raw, axis=1, keepdims=True)
    exp_raw = np.exp(raw - raw_max)
    probs = exp_raw / np.sum(exp_raw, axis=1, keepdims=True)

    # Argmax per timestep
    best_indices = np.argmax(probs, axis=1)  # [T]
    best_scores = probs[np.arange(seq_len), best_indices]  # [T]

    # Collapse consecutive duplicates (retain the timestep each run began at so
    # the rescorer can read a confusable's softmax mass at that aligned step).
    collapsed = []
    collapsed_scores = []
    collapsed_timesteps = []
    for i, idx in enumerate(best_indices):
        if i == 0 or idx != best_indices[i - 1]:
            collapsed.append(int(idx))
            collapsed_scores.append(float(best_scores[i]))
            collapsed_timesteps.append(i)

    # Remove blank (index 0) and join
    chars = []
    scores = []
    char_timesteps = []
    for idx, score, t in zip(collapsed, collapsed_scores, collapsed_timesteps):
        if idx != 0:
            if idx < len(charset):
                chars.append(charset[idx])
            else:
                chars.append(chr(idx))
            scores.append(score)
            char_timesteps.append(t)

    text = "".join(chars)
    confidence = float(np.mean(scores)) if scores else 0.0

    # Plausibility gate — suppress the degenerate near-all-blank collapse so
    # low-resolution noise is never stored as a plate label. Count ALPHANUMERIC
    # characters only ("DW RO 309" → 7): the charset carries symbol glyphs
    # ('[', ')', '.', …) that a blurred crop hallucinates, and they must not
    # count toward "long enough to be a plate".
    min_len = int(config.get("minTextLength", DEFAULT_MIN_TEXT_LENGTH))
    min_conf = float(config.get("minTextConfidence", DEFAULT_MIN_TEXT_CONFIDENCE))
    region = str(config.get("plateRegion", ABSENT_PLATE_REGION))
    alnum_count = sum(1 for c in text if c.isalnum())
    # Symbol-blind gate: when a plate grammar/region is ACTIVE, no real plate
    # contains a non-alphanumeric glyph (spaces allowed — word gaps are real).
    # Observed live labels "[miv4o)" / "6ou.2" (2026-07-22) sailed through the
    # length-only gate; reject them exactly like the other implausible reads.
    has_symbol_glyph = region != "off" and any(
        not (c.isalnum() or c.isspace()) for c in text
    )
    # …except the SEAL. A German plate physically carries a registration seal
    # between the district code and the letter group, and the recognizer
    # transcribes that badge as a punctuation glyph ('.', ':', '-', …). Once the
    # quad warp (2026-08-20) started delivering legible strips, the correct
    # decodes arrived as "DN.D7578" / "DN-P 1091" / "DN:HF 974" and this gate
    # threw away every one of them — the warp scored 1/21 instead of 4/21 on the
    # cam-617 ground-truth set purely because of this line.
    #
    # The exemption is deliberately narrow: EXACTLY ONE seal-like glyph, strictly
    # INTERIOR, and the remainder must CONFORM to the region grammar. Noise
    # cannot be laundered by it — the 2026-07-22 junk ("6ou.2x", "[miv4o)") has
    # either the wrong count, the wrong position, or no conforming remainder.
    if has_symbol_glyph:
        repaired = _strip_seal_glyph(text, region)
        if repaired is not None:
            return {
                "kind": "text",
                "text": repaired,
                "confidence": round(confidence, 4),
                "formatValid": True,
            }
    if alnum_count < min_len or confidence < min_conf or has_symbol_glyph:
        return {
            "kind": "text",
            "text": "",
            "confidence": round(confidence, 4),
        }

    result = {
        "kind": "text",
        "text": text,
        "confidence": round(confidence, 4),
    }

    # Post-hoc format validation + confusable rescoring. 'off' disables it
    # entirely (no grammar lookup, no formatValid annotation).
    if region != "off":
        pattern_src = PLATE_GRAMMARS.get(region)
        if pattern_src is not None:
            pattern = re.compile(pattern_src)
            # Grammar + confusables run on the NORMALISED read (upper, alnum
            # only). Reached only past the plausibility gate above, so a repair
            # can never turn sub-legible noise into a more credible string: the
            # symbol gate already rejected the reads that would qualify, and the
            # normalised length equals the alnum count the length gate cleared.
            norm_text, norm_scores, norm_timesteps = _normalize_read(
                chars, scores, char_timesteps
            )
            fixed_text, format_valid = _rescore_for_grammar(
                norm_text, norm_scores, norm_timesteps, probs, list(charset), pattern
            )
            result["text"] = fixed_text
            result["formatValid"] = format_valid

    return result