"""S5 — CTC forced alignment of a transcript to character CTC emissions (Viterbi).

Given log-softmax emissions ``[num_frames, vocab]`` from a character CTC model (e.g.
``mpoyraz_ctc_tr``) and the primary transcript, find the single most likely CTC path that
emits exactly the transcript's characters (words separated by the model's word-delimiter
token) and read word intervals off that path. Plain numpy, no torchaudio dependency
(``torchaudio.functional.forced_align`` is deprecated in the 2.8 line, R8).

Text handling: each transcript token is NFC + ``tr_lower``-ed, circumflexes are folded
(â/î/û → a/i/u: the model has no such letters) and characters outside the vocabulary that
are punctuation are dropped for alignment only. A token with no alignable character
(``…``, ``-``) is glued to the previous word (or the next one at the start), so the
returned words keep their original spelling and punctuation. Any other character
outside the vocabulary (digits, foreign letters) makes the segment unalignable.
"""

from __future__ import annotations

import itertools
import math
import unicodedata
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from nida.data.types import TimedWord
from nida.text.turkish import nfc, tr_lower

if TYPE_CHECKING:
    import numpy as np
    from numpy.typing import NDArray

__all__ = ["AlignmentResult", "align_words", "alignment_units"]

_CIRCUMFLEX_FOLD = str.maketrans({"â": "a", "î": "i", "û": "u"})
# Typographic apostrophes (U+2019, U+02BC) are the same letter-joining apostrophe as U+0027.
_APOSTROPHES = str.maketrans({"\u2019": "'", "\u02bc": "'"})
#: Allowed difference between ``num_frames * hop`` and the segment duration, in frames
#: (wav2vec2 yields ``floor((samples - 400) / 320) + 1`` frames, i.e. up to ~1 frame short).
FRAME_SLACK = 2


@dataclass(frozen=True)
class AlignmentResult:
    """Either ``words`` (aligned) or ``failure`` (a machine-readable reason) is set."""

    words: tuple[TimedWord, ...] | None
    failure: str | None
    details: Mapping[str, object] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if (self.words is None) == (self.failure is None):
            raise ValueError("AlignmentResult needs exactly one of words / failure")
        if self.words is not None and self.details:
            raise ValueError("an aligned result carries no failure details")


def alignment_units(
    tokens: Sequence[str], vocab: Mapping[str, int]
) -> tuple[list[str], list[str]] | str:
    """Group transcript tokens into alignable words.

    Returns:
        ``(display_words, alignment_strings)`` of equal length, or a failure string
        ``"unalignable_chars:<chars>"`` / ``"no_alignable_text"``. Apostrophes join
        letters only: leading/trailing ones are not aligned, and a token made only of
        punctuation (``…``, ``—``, ``'``) is glued to the previous word (or the next one at
        the start) *without* the space between them (``a — b`` → ``a—``, ``b``), because a
        returned word cannot contain whitespace.
    """
    display: list[str] = []
    targets: list[str] = []
    pending_prefix = ""
    bad: set[str] = set()
    for raw in tokens:
        token = nfc(raw)
        if not token or any(ch.isspace() for ch in token):
            raise ValueError(f"invalid transcript token {raw!r}")
        chars: list[str] = []
        for ch in tr_lower(token).translate(_CIRCUMFLEX_FOLD).translate(_APOSTROPHES):
            if ch in vocab:
                chars.append(ch)
            elif unicodedata.category(ch).startswith("P"):
                continue
            else:
                bad.add(ch)
        chars = list("".join(chars).strip("'"))
        if not chars:
            if display:
                display[-1] += token
            else:
                pending_prefix += token
            continue
        display.append(pending_prefix + token)
        pending_prefix = ""
        targets.append("".join(chars))
    if bad:
        return "unalignable_chars:" + "".join(sorted(bad))
    if not targets:
        return "no_alignable_text"
    return display, targets


def _viterbi(log_probs: NDArray[np.float64], labels: Sequence[int], blank: int) -> list[int] | None:
    """Best CTC state path; ``None`` if the label sequence cannot fit in the frames.

    States are the blank-extended sequence ``b l1 b l2 b … lN b`` (index ``2k+1`` is
    label ``k``). Returns the state index for every frame.
    """
    import numpy as np

    num_frames = log_probs.shape[0]
    ext = np.full(2 * len(labels) + 1, blank, dtype=np.int64)
    ext[1::2] = labels
    num_states = ext.shape[0]
    # A skip s-2 -> s is allowed into a label that differs from the label two back.
    can_skip = np.zeros(num_states, dtype=bool)
    can_skip[3::2] = ext[3::2] != ext[1:-2:2]
    neg_inf = -np.inf
    score = np.full(num_states, neg_inf)
    score[0] = log_probs[0, ext[0]]
    if num_states > 1:
        score[1] = log_probs[0, ext[1]]
    back = np.zeros((num_frames, num_states), dtype=np.int8)  # 0 stay, 1 step, 2 skip
    for t in range(1, num_frames):
        stay = score
        step = np.concatenate(([neg_inf], score[:-1]))
        skip = np.concatenate(([neg_inf, neg_inf], score[:-2]))
        skip = np.where(can_skip, skip, neg_inf)
        stacked = np.stack([stay, step, skip])
        choice = np.argmax(stacked, axis=0)
        best = np.take_along_axis(stacked, choice[None, :], axis=0)[0]
        score = best + log_probs[t, ext]
        back[t] = choice
    final_candidates = [num_states - 1] + ([num_states - 2] if num_states > 1 else [])
    end = max(final_candidates, key=lambda s: score[s])
    if not math.isfinite(float(score[end])):
        return None
    path = [0] * num_frames
    state = end
    for t in range(num_frames - 1, -1, -1):
        path[t] = state
        state -= int(back[t, state])
    return path


def align_words(
    tokens: Sequence[str],
    log_probs: NDArray[np.float32],
    vocab: Mapping[str, int],
    *,
    blank_token: str,
    word_delimiter: str,
    frame_hop_ms: float,
    duration_ms: int,
) -> AlignmentResult:
    """Align transcript ``tokens`` (whitespace-split words) to CTC ``log_probs``.

    Args:
        tokens: Transcript words in order.
        log_probs: ``[num_frames, vocab_size]`` emissions. Log-softmax is re-applied
            here (idempotent on log-probabilities), so raw logits cannot inflate
            ``align_score``. ``-inf`` entries are allowed; NaN and ``+inf`` are not.
        vocab: Token → id (characters, ``blank_token``, ``word_delimiter``).
        blank_token: CTC blank (``[PAD]`` for wav2vec2 HF models).
        word_delimiter: Word separator token (``|``).
        frame_hop_ms: Frame stride in ms (20 for wav2vec2).
        duration_ms: Segment length. ``num_frames * frame_hop_ms`` must be within
            :data:`FRAME_SLACK` frames of it (catches unresampled audio); word ends are
            clamped to it.

    Returns:
        :class:`AlignmentResult` with words carrying ``align_score`` = mean emission
        probability of the word's character frames, or a failure reason
        (``unalignable_chars``, ``no_alignable_text``, ``too_few_frames``, ``no_path``,
        ``words_beyond_audio``).

    Raises:
        ValueError: malformed inputs (shape, vocab ids not ``0..V-1``, NaN/+inf, bad hop,
            frame count inconsistent with the duration).
    """
    import numpy as np

    if log_probs.ndim != 2 or log_probs.shape[1] != len(vocab):
        raise ValueError(f"log_probs shape {log_probs.shape} does not match vocab {len(vocab)}")
    if sorted(vocab.values()) != list(range(len(vocab))):
        raise ValueError("vocab ids must be exactly 0..V-1")
    if np.any(np.isnan(log_probs)) or np.any(np.isposinf(log_probs)):
        raise ValueError("log_probs contain NaN or +inf")
    if blank_token not in vocab or word_delimiter not in vocab:
        raise ValueError("blank or word delimiter missing from vocab")
    if not (math.isfinite(frame_hop_ms) and frame_hop_ms >= 1):
        raise ValueError(f"frame_hop_ms must be a finite number >= 1 ms, got {frame_hop_ms!r}")
    if type(duration_ms) is not int or duration_ms <= 0:
        raise ValueError(f"invalid duration_ms {duration_ms!r}")
    if abs(log_probs.shape[0] * frame_hop_ms - duration_ms) > FRAME_SLACK * frame_hop_ms:
        raise ValueError(
            f"{log_probs.shape[0]} frames x {frame_hop_ms} ms do not match {duration_ms} ms"
        )
    with np.errstate(invalid="ignore"):
        lp = log_probs.astype(np.float64)
        row_max = np.max(lp, axis=1, keepdims=True)
        if not np.all(np.isfinite(row_max)):
            raise ValueError("a frame has no finite log-probability")
        lp = lp - (row_max + np.log(np.sum(np.exp(lp - row_max), axis=1, keepdims=True)))
    special = {blank_token, word_delimiter}
    char_vocab = {tok: idx for tok, idx in vocab.items() if len(tok) == 1 and tok not in special}
    units = alignment_units(tokens, char_vocab)
    if isinstance(units, str):
        reason, _, chars = units.partition(":")
        return AlignmentResult(None, reason, {"chars": chars} if chars else {})
    display, targets = units

    labels: list[int] = []
    owner: list[int] = []  # word index for each label, -1 for delimiters
    for index, target in enumerate(targets):
        if index:
            labels.append(vocab[word_delimiter])
            owner.append(-1)
        for ch in target:
            labels.append(char_vocab[ch])
            owner.append(index)
    repeats = sum(1 for a, b in itertools.pairwise(labels) if a == b)
    if log_probs.shape[0] < len(labels) + repeats:
        return AlignmentResult(
            None,
            "too_few_frames",
            {"frames": int(log_probs.shape[0]), "needed": len(labels) + repeats},
        )
    path = _viterbi(lp, labels, vocab[blank_token])
    if path is None:
        return AlignmentResult(None, "no_path")

    first: list[int | None] = [None] * len(targets)
    last: list[int | None] = [None] * len(targets)
    probs: list[list[float]] = [[] for _ in targets]
    for frame, state in enumerate(path):
        if state % 2 == 0:
            continue
        word = owner[state // 2]
        if word < 0:
            continue
        if first[word] is None:
            first[word] = frame
        last[word] = frame
        probs[word].append(math.exp(float(lp[frame, labels[state // 2]])))
    words: list[TimedWord] = []
    for index, text in enumerate(display):
        start_frame, end_frame = first[index], last[index]
        if start_frame is None or end_frame is None:
            return AlignmentResult(None, "no_path", {"word_index": index})
        start_ms = min(round(start_frame * frame_hop_ms), duration_ms - 1)
        end_ms = min(round((end_frame + 1) * frame_hop_ms), duration_ms)
        end_ms = max(end_ms, start_ms + 1)
        if words and start_ms < words[-1].end_ms:
            # Only reachable through the end clamp (emissions slightly longer than the
            # audio): words squeezed past the end cannot be timed, never let them overlap.
            return AlignmentResult(None, "words_beyond_audio", {"word_index": index})
        score = sum(probs[index]) / len(probs[index])
        words.append(TimedWord(text, start_ms, end_ms, score))
    return AlignmentResult(tuple(words), None)