scoring.py

"""
Backend-agnostic scoring. Takes a Transcript from any ASR backend plus the
reference passage, returns metrics + an error table.

Nothing in here touches a GPU or the network, which is why the app can run on a
free CPU Space: only transcription needs compute.
"""

from __future__ import annotations

import re
import string
import unicodedata
from functools import lru_cache
from typing import Iterable, Sequence

import pandas as pd

from asr_backends import Transcript, is_nan

# ---------------------------------------------------------------------------
# Optional deps
# ---------------------------------------------------------------------------

try:
    from Levenshtein import distance as _lev
    from Levenshtein import ratio as _ratio
except ImportError:
    import difflib

    def _lev(a: str, b: str) -> int:
        sm = difflib.SequenceMatcher(None, a, b)
        n = max(len(a), len(b))
        return n - int(sm.ratio() * n)

    def _ratio(a: str, b: str) -> float:
        return difflib.SequenceMatcher(None, a, b).ratio()

try:
    from num2words import num2words
except ImportError:
    num2words = None


# ---------------------------------------------------------------------------
# Normalisation -- most "errors" in a naive implementation are encoding noise
# ---------------------------------------------------------------------------

ZERO_WIDTH = dict.fromkeys(map(ord, "\u200b\u200c\u200d\ufeff"), None)
DEVANAGARI_DIGITS = {ord(c): str(i) for i, c in enumerate("०१२३४५६७८९")}
DEV_PUNCT = "।॥"
MATRAS = re.compile(r"[\u0900-\u0903\u093A-\u094F\u0951-\u0957\u0962\u0963]")

CONTRACTIONS = {
    "cant": "cannot", "dont": "do not", "wont": "will not", "im": "i am",
    "ive": "i have", "id": "i would", "ill": "i will", "its": "it is",
    "lets": "let us", "thats": "that is", "youre": "you are", "hes": "he is",
    "shes": "she is", "theyre": "they are", "isnt": "is not", "arent": "are not",
    "wasnt": "was not", "didnt": "did not", "doesnt": "does not",
    "couldnt": "could not", "wouldnt": "would not", "shouldnt": "should not",
}


def _expand_numbers(text: str, lang: str) -> str:
    if num2words is None:
        return text

    def sub(m: re.Match) -> str:
        try:
            return num2words(int(m.group()), lang="hi" if lang == "hi" else "en")
        except Exception:
            return m.group()

    return re.sub(r"\d+", sub, text)


def normalise(text: str, lang: str) -> list[str]:
    """Canonical token list. Order of operations matters."""
    text = unicodedata.normalize("NFC", text)      # unifies the two encodings of ड़
    text = text.translate(ZERO_WIDTH)

    if lang == "hi":
        text = text.translate(DEVANAGARI_DIGITS)
        text = text.replace("ॐ", "ओम")
        text = re.sub(f"[{DEV_PUNCT}]", " ", text)
        text = text.replace("ँ", "ं")               # chandrabindu ~ anusvara
    else:
        text = text.lower().replace("\u2019", "'").replace("-", " ")

    text = _expand_numbers(text, lang)
    text = text.translate(str.maketrans("", "", string.punctuation))
    tokens = text.split()

    if lang == "en":
        tokens = [CONTRACTIONS.get(t, t) for t in tokens]
        tokens = [w for t in tokens for w in t.split()]
    return tokens


def skeleton(word: str, lang: str) -> str:
    """Vowel-stripped form: equal skeletons mean same consonants, wrong vowels."""
    if lang == "hi":
        return MATRAS.sub("", word)
    return re.sub(r"[aeiou]", "", word) or word


@lru_cache(maxsize=200_000)
def similarity(a: str, b: str) -> float:
    return _ratio(a, b)


# ---------------------------------------------------------------------------
# Alignment: Needleman-Wunsch weighted by edit distance
# ---------------------------------------------------------------------------
# difflib only matches byte-identical tokens, so बिगडा vs बिगड़ा becomes a
# delete + insert and the two words are never compared to each other.

GAP_COST = 0.62     # < 1.0 so a near-match always beats delete + insert


def align(ref: Sequence[str], hyp: Sequence[str]) -> list[tuple[str | None, str | None]]:
    n, m = len(ref), len(hyp)
    dist = [[0.0] * (m + 1) for _ in range(n + 1)]
    back = [[""] * (m + 1) for _ in range(n + 1)]

    for i in range(1, n + 1):
        dist[i][0], back[i][0] = i * GAP_COST, "D"
    for j in range(1, m + 1):
        dist[0][j], back[0][j] = j * GAP_COST, "I"

    for i in range(1, n + 1):
        ri = ref[i - 1]
        for j in range(1, m + 1):
            sub = dist[i - 1][j - 1] + (1.0 - similarity(ri, hyp[j - 1]))
            dele = dist[i - 1][j] + GAP_COST
            ins = dist[i][j - 1] + GAP_COST
            best = min(sub, dele, ins)
            dist[i][j] = best
            back[i][j] = "M" if best == sub else ("D" if best == dele else "I")

    pairs: list[tuple[str | None, str | None]] = []
    i, j = n, m
    while i > 0 or j > 0:
        op = back[i][j] if (i and j) else ("D" if i else "I")
        if op == "M":
            pairs.append((ref[i - 1], hyp[j - 1])); i -= 1; j -= 1
        elif op == "D":
            pairs.append((ref[i - 1], None)); i -= 1
        else:
            pairs.append((None, hyp[j - 1])); j -= 1
    pairs.reverse()
    return pairs


# ---------------------------------------------------------------------------
# Error taxonomy
# ---------------------------------------------------------------------------

SIMILAR_HI = [set("बवभ"), set("सशष"), set("दध"), set("तट"), set("कख"), set("गघ"),
              set("जझ"), set("पफ"), set("नण"), set("रड़"), set("लर")]
SIMILAR_EN = [set("bvp"), set("sz"), set("td"), set("kg"), set("fp"), set("lr"),
              set("mn"), set("jy")]

LABELS = {
    "extra":         ("अतिरिक्त शब्द", "Extra word"),
    "omission":      ("छूटा हुआ शब्द", "Omitted word"),
    "vowel":         ("मात्रा दोष", "Vowel error"),
    "phonetic":      ("ध्वनि भ्रम", "Confusable sound"),
    "pronunciation": ("उच्चारण दोष", "Mispronounced"),
    "order":         ("अक्षर क्रम", "Letter order"),
    "substitution":  ("गलत शब्द", "Wrong word"),
}

SEVERITY = {"ok": 0, "vowel": 1, "phonetic": 1, "pronunciation": 2, "order": 2,
            "omission": 3, "extra": 3, "substitution": 4}


def _label(code: str, lang: str) -> str:
    hi, en = LABELS[code]
    return f"{hi} / {en}" if lang == "hi" else en


def classify(ref: str | None, hyp: str | None, lang: str) -> str:
    if ref is None:
        return "extra"
    if hyp is None:
        return "omission"
    if ref == hyp:
        return "ok"

    ed = _lev(ref, hyp)
    if skeleton(ref, lang) == skeleton(hyp, lang):
        return "vowel"

    groups = SIMILAR_HI if lang == "hi" else SIMILAR_EN
    if ed <= 2 and any((set(ref) & g) and (set(hyp) & g) for g in groups):
        return "phonetic"
    if similarity(ref, hyp) >= 0.75 or ed <= 2:
        return "pronunciation"
    if sorted(ref) == sorted(hyp):
        return "order"
    return "substitution"


# ---------------------------------------------------------------------------
# Metrics
# ---------------------------------------------------------------------------

def cer(ref_tokens: Iterable[str], hyp_tokens: Iterable[str]) -> float:
    r, h = " ".join(ref_tokens), " ".join(hyp_tokens)
    return _lev(r, h) / max(1, len(r))


def score(expected: str, tr: Transcript, lang: str) -> tuple[dict, pd.DataFrame]:
    ref = normalise(expected, lang)
    hyp = normalise(tr.text, lang)

    if not ref:
        return {"error": "The passage is empty."}, pd.DataFrame()

    pairs = align(ref, hyp)

    # map normalised hypothesis token -> ASR confidence, when the backend has it
    conf: dict[str, float] = {}
    if tr.has_confidence:
        for w in tr.words:
            toks = normalise(w.text, lang)
            if toks:
                conf.setdefault(toks[0], w.prob)

    rows, sub, dele, ins, soft = [], 0, 0, 0, 0.0
    for r, h in pairs:
        code = classify(r, h, lang)
        if code == "ok":
            soft += 1.0
            continue
        if code == "extra":
            ins += 1
        elif code == "omission":
            dele += 1
        else:
            sub += 1
            soft += similarity(r, h)          # partial credit for a near miss

        row = {
            "अपेक्षित / Expected": r or "",
            "सुना गया / Heard": h or "",
            "प्रकार / Error type": _label(code, lang),
            "समानता / Similarity": round(similarity(r or "", h or ""), 2),
        }
        if tr.has_confidence:
            c = conf.get(h) if h else None
            row["ASR conf."] = None if (c is None or is_nan(c)) else round(c, 2)
        row["_sev"] = SEVERITY[code]
        rows.append(row)

    n = len(ref)
    wer = (sub + dele + ins) / n
    exact = 100.0 * max(0, n - sub - dele) / n
    lenient = 100.0 * soft / n

    dur = tr.speech_seconds or (
        tr.words[-1].end - tr.words[0].start if len(tr.words) > 1 else 0.0)
    wpm = round(60.0 * len(hyp) / dur, 1) if dur > 0.5 else None

    pauses = sum(1 for a, b in zip(tr.words, tr.words[1:]) if b.start - a.end > 0.7)

    metrics = {
        "📝 Transcribed": tr.text,
        "✅ Word accuracy (%)": round(exact, 2),
        "🎯 Lenient score (%)": round(lenient, 2),
        "📉 WER (%)": round(100 * wer, 2),
        "🔤 CER (%)": round(100 * cer(ref, hyp), 2),
        "⏱️ Speaking rate (wpm)": wpm,
        "⏸️ Long pauses (>0.7s)": pauses if tr.words else "n/a",
        "🔢 Errors": {"substitutions": sub, "omissions": dele, "insertions": ins},
        "⚙️ Backend": f"{tr.backend}:{tr.model} ({tr.latency_s}s)",
    }
    if tr.has_confidence:
        unclear = [w.text for w in tr.words if not is_nan(w.prob) and w.prob < 0.45]
        metrics["🤔 Unclear words"] = unclear[:10] or "—"

    df = pd.DataFrame(rows)
    if not df.empty:
        df = (df.sort_values("_sev", ascending=False)
                .drop(columns="_sev")
                .reset_index(drop=True))
    return metrics, df