reduction/stt.py

from __future__ import annotations

import difflib
import re
from dataclasses import dataclass

from faster_whisper import WhisperModel


_STRIP_RE = re.compile(r"^[\s\.,!?;:\"'„“”‚‘’\(\)\[\]…—–-]+|[\s\.,!?;:\"'„“”‚‘’\(\)\[\]…—–-]+$")


@dataclass
class Word:
    text: str       # normalized form for TTS (lowercased, punctuation stripped)
    raw: str        # original Whisper token (kept for debugging / display)
    start: float    # seconds in original audio
    end: float
    probability: float = float("nan")  # Whisper per-word probability (NaN if unknown)

    @property
    def duration(self) -> float:
        return self.end - self.start


def _normalize(token: str) -> str:
    return _STRIP_RE.sub("", token).lower()


def transcribe_words(
    audio_path: str,
    model: WhisperModel,
    language: str = "cs",
) -> tuple[list[Word], str]:
    """Transcribe `audio_path` and return word-level segmentation + full text.

    Empty / pure-punctuation tokens are skipped.
    """
    segments, _info = model.transcribe(
        audio_path, language=language, word_timestamps=True
    )
    words: list[Word] = []
    full_text: list[str] = []
    for seg in segments:
        full_text.append(seg.text)
        if not seg.words:
            continue
        for w in seg.words:
            normalized = _normalize(w.word)
            if not normalized:
                continue
            prob = getattr(w, "probability", float("nan"))
            words.append(
                Word(
                    text=normalized,
                    raw=w.word,
                    start=w.start,
                    end=w.end,
                    probability=float(prob) if prob is not None else float("nan"),
                )
            )
    return words, "".join(full_text).strip()


def load_model(
    size: str = "large-v3",
    device: str = "cpu",
    compute_type: str = "int8",
    download_root: str = "models/whisper",
) -> WhisperModel:
    return WhisperModel(
        size, device=device, compute_type=compute_type, download_root=download_root
    )


def tokenize_transcript(text: str) -> list[str]:
    """Split a free-form transcript into normalized word tokens (lowercase, no punct)."""
    return [t for t in (_normalize(tok) for tok in text.split()) if t]


def _char_dist(a: str, b: str) -> float:
    """Normalized character-level edit distance in [0, 1]. 0 = identical."""
    if not a and not b:
        return 0.0
    return 1.0 - difflib.SequenceMatcher(a=a, b=b, autojunk=False).ratio()


def align_to_groundtruth(
    whisper_words: list[Word], gt_tokens: list[str], gap_penalty: float = 0.7
) -> list[Word]:
    """Replace Whisper word identities with ground-truth tokens, keeping Whisper timing.

    Uses Needleman–Wunsch alignment with character-level edit distance as substitution cost.
    Behavior:
      - Aligned (W_i, GT_j): emit Word(text=GT_j, start/end from W_i).
      - GT_j unmatched: try to interpolate timing from neighboring aligned tokens; if
        impossible (start/end of utterance with no anchor), drop with a stderr warning.
      - W_i unmatched: drop (Whisper inserted a phantom token).
    """
    if not whisper_words:
        return []
    if not gt_tokens:
        return list(whisper_words)

    n, m = len(whisper_words), len(gt_tokens)
    dp = [[0.0] * (m + 1) for _ in range(n + 1)]
    bt = [[""] * (m + 1) for _ in range(n + 1)]
    for i in range(1, n + 1):
        dp[i][0] = i * gap_penalty
        bt[i][0] = "D"
    for j in range(1, m + 1):
        dp[0][j] = j * gap_penalty
        bt[0][j] = "I"
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            sub = dp[i - 1][j - 1] + _char_dist(whisper_words[i - 1].text, gt_tokens[j - 1])
            dele = dp[i - 1][j] + gap_penalty
            ins = dp[i][j - 1] + gap_penalty
            best = min(sub, dele, ins)
            dp[i][j] = best
            bt[i][j] = "M" if best == sub else ("D" if best == dele else "I")

    # backtrack
    pairs: list[tuple[int | None, int | None]] = []
    i, j = n, m
    while i > 0 or j > 0:
        op = bt[i][j]
        if op == "M":
            pairs.append((i - 1, j - 1))
            i -= 1
            j -= 1
        elif op == "D":
            pairs.append((i - 1, None))
            i -= 1
        else:
            pairs.append((None, j - 1))
            j -= 1
    pairs.reverse()

    # First pass: emit matched + skip Whisper-only; collect GT-only with neighbor anchors.
    out: list[Word | None] = []
    pending_gt: list[int] = []  # GT indices waiting for next anchor
    last_end: float | None = None
    for w_idx, g_idx in pairs:
        if w_idx is not None and g_idx is not None:
            ww = whisper_words[w_idx]
            # if there are pending GT tokens, distribute them between last_end and ww.start
            if pending_gt:
                if last_end is not None:
                    span_start, span_end = last_end, ww.start
                    n_pend = len(pending_gt)
                    for k, gi in enumerate(pending_gt):
                        ts = span_start + (span_end - span_start) * k / n_pend
                        te = span_start + (span_end - span_start) * (k + 1) / n_pend
                        out.append(Word(text=gt_tokens[gi], raw=gt_tokens[gi], start=ts, end=te))
                else:
                    # at start of utterance with no anchor — assign 0..ww.start span
                    n_pend = len(pending_gt)
                    for k, gi in enumerate(pending_gt):
                        ts = ww.start * k / n_pend
                        te = ww.start * (k + 1) / n_pend
                        out.append(Word(text=gt_tokens[gi], raw=gt_tokens[gi], start=ts, end=te))
                pending_gt = []
            out.append(Word(text=gt_tokens[g_idx], raw=gt_tokens[g_idx], start=ww.start, end=ww.end))
            last_end = ww.end
        elif w_idx is not None and g_idx is None:
            # Whisper-only — skip, but advance time anchor
            last_end = whisper_words[w_idx].end
        else:  # GT-only
            pending_gt.append(g_idx)  # type: ignore[arg-type]

    # Trailing pending GT tokens (no right-anchor): assign small constant duration
    if pending_gt:
        if last_end is not None:
            span_start = last_end
            span_end = last_end + 0.3 * len(pending_gt)  # rough fallback
            n_pend = len(pending_gt)
            for k, gi in enumerate(pending_gt):
                ts = span_start + (span_end - span_start) * k / n_pend
                te = span_start + (span_end - span_start) * (k + 1) / n_pend
                out.append(Word(text=gt_tokens[gi], raw=gt_tokens[gi], start=ts, end=te))

    return [w for w in out if w is not None]