markers.py

"""Extract language attrition markers from a whisper-timestamped transcript.

Covers the original 13 Schmid-style + temporal + lexical markers, plus
Phase-1 acoustic markers (F0, voice quality, spectral tilt, MFCCs).
"""

from __future__ import annotations

import re
import statistics
from collections import Counter

import pyphen
from lexical_diversity import lex_div

from acoustic_markers import extract_acoustic_markers
from alignment_markers import extract_alignment_markers

# ---------------------------------------------------------------------------
# Word lists
# ---------------------------------------------------------------------------

SPANISH_FILLERS = {
    "eh", "ah", "um", "uh", "este", "em", "mmm", "pues", "bueno",
    "mm", "hm", "ehm", "uhm",
}

# Bracket tokens for filled pauses: CrisperWhisper [UH]/[UM] or
# whisper-timestamped [*] disfluency markers
_BRACKET_FP = re.compile(r"^\[U[HM]\]$", re.IGNORECASE)
_DISFLUENCY_MARKER = re.compile(r"^\[\*\]$")

SPANISH_ARTICLES = {"el", "la", "los", "las", "un", "una", "unos", "unas"}
SPANISH_PREPOSITIONS = {
    "de", "en", "a", "con", "por", "para", "sin", "sobre",
    "entre", "hacia", "hasta", "desde", "según", "durante",
}
SPANISH_PRONOUNS = {
    "yo", "tú", "él", "ella", "nosotros", "nosotras",
    "ellos", "ellas", "usted", "ustedes",
    "me", "te", "se", "nos", "le", "les", "lo", "la",
    "mí", "ti", "sí", "conmigo", "contigo", "consigo",
    "que", "quien", "cual", "donde", "como",
    "este", "ese", "aquel", "esto", "eso", "aquello",
}

FUNCTION_WORD_SETS: dict[str, set[str]] = {
    "article": SPANISH_ARTICLES,
    "preposition": SPANISH_PREPOSITIONS,
    "pronoun": SPANISH_PRONOUNS,
}

# Common Spanish verbs (high-frequency) for word-class tagging
SPANISH_VERBS = {
    "es", "era", "fue", "ser", "estar", "está", "estaba", "estuvo",
    "tiene", "tenía", "tuvo", "tener", "hacer", "hace", "hizo",
    "ir", "va", "iba", "fue", "viene", "venir", "dar", "dio",
    "decir", "dice", "dijo", "saber", "sabe", "sabía",
    "poder", "puede", "podía", "pudo", "querer", "quiere", "quería",
    "ver", "ve", "vio", "veía", "poner", "pone", "puso",
    "salir", "sale", "salió", "llegar", "llega", "llegó",
    "pasar", "pasa", "pasó", "quedar", "queda", "quedó",
    "creer", "cree", "pensaba", "pensar", "piensa",
    "hay", "había", "haber", "han", "ha",
    "entrar", "entra", "entró", "caer", "cae", "cayó",
    "llevar", "lleva", "llevó", "seguir", "sigue", "siguió",
    "encontrar", "encuentra", "encontró",
    "empezar", "empieza", "empezó", "intentar", "intenta", "intentó",
    "sentir", "siente", "sintió",
}

# German words unlikely to appear in Spanish — for code-switching detection
GERMAN_WORDS = {
    "und", "aber", "auch", "schon", "jetzt", "dann", "nicht", "oder",
    "weil", "dass", "noch", "hier", "dort", "halt", "genau",
    "eigentlich", "vielleicht", "natürlich", "wirklich", "immer",
    "wieder", "heute", "morgen", "gestern", "gerade", "trotzdem",
    "deswegen", "deshalb", "allerdings", "übrigens", "sozusagen",
    "naja", "eben", "doch", "mal", "nur", "sehr", "ganz", "etwas",
    "alles", "nichts", "niemand", "jemand", "etwas", "irgendwie",
    "auf", "aus", "bei", "mit", "nach", "seit", "von", "zu",
    "über", "unter", "vor", "hinter", "neben", "zwischen",
    "ich", "du", "wir", "ihr", "sie", "mein", "dein", "sein",
    "kein", "diese", "jede", "welche", "solche",
    "können", "müssen", "sollen", "wollen", "dürfen", "mögen",
    "machen", "gehen", "kommen", "sagen", "wissen", "sehen",
    "geben", "nehmen", "finden", "denken", "stehen", "lassen",
    "sprechen", "halten", "bringen", "leben", "fahren", "arbeiten",
    "spielen", "laufen", "schreiben", "lesen", "essen", "trinken",
    "schlafen", "kaufen", "verkaufen", "helfen", "lernen",
    "arbeit", "schule", "kindergarten", "ausbildung", "beruf",
    "wohnung", "straße", "stadt", "kirche",
    "ja", "nein", "bitte", "danke", "stimmt", "klar", "richtig",
}

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_pyphen_dic = pyphen.Pyphen(lang="es_ES")


def _clean_word(w: str) -> str:
    """Lowercase and strip punctuation for comparison."""
    return re.sub(r"[^\w]", "", w.lower())


def _is_disfluency_marker(word: str) -> bool:
    """Check if word is a whisper-timestamped [*] disfluency marker."""
    return bool(_DISFLUENCY_MARKER.match(word.strip()))


def _is_filler(word: str) -> bool:
    """Check if word is a filled pause (bracket token, [*] marker, or Spanish filler)."""
    w = word.strip()
    if _BRACKET_FP.match(w):
        return True
    if _is_disfluency_marker(w):
        return True
    return _clean_word(w) in SPANISH_FILLERS


def _clean_tokens(chunks) -> list[str]:
    """Return lowered, punctuation-stripped tokens excluding fillers and [*] markers."""
    out = []
    for c in chunks:
        if _is_disfluency_marker(c["text"]):
            continue
        w = _clean_word(c["text"])
        if w and not _is_filler(c["text"]):
            out.append(w)
    return out


def _syllable_count(word: str) -> int:
    """Count syllables using pyphen Spanish dictionary."""
    parts = _pyphen_dic.inserted(word).split("-")
    return max(len(parts), 1)


def _word_class(word: str) -> str:
    """Simple rule-based word-class tagger for Spanish."""
    w = _clean_word(word)
    if not w:
        return "other"
    if w in SPANISH_ARTICLES:
        return "article"
    if w in SPANISH_PREPOSITIONS:
        return "preposition"
    if w in SPANISH_PRONOUNS:
        return "pronoun"
    if w in SPANISH_VERBS:
        return "verb"
    # heuristic: treat remaining content words as nouns (open class)
    return "noun"


def _function_word_category(word: str) -> str | None:
    """Return function-word category or None if content word."""
    w = _clean_word(word)
    for cat, wordset in FUNCTION_WORD_SETS.items():
        if w in wordset:
            return cat
    return None


# ---------------------------------------------------------------------------
# Marker extractors
# ---------------------------------------------------------------------------

def _filled_pauses(chunks) -> list[dict]:
    """SDM-1: Filled pauses (Spanish fillers + [*] disfluency markers)."""
    fps = []
    for i, c in enumerate(chunks):
        if _is_disfluency_marker(c["text"]):
            fps.append({
                "index": i,
                "text": "[*]",
                "source": "whisper-disfluency",
                **c["timestamp"],
            })
        elif _is_filler(c["text"]):
            fps.append({
                "index": i,
                "text": c["text"].strip(),
                "source": "filler-word",
                **c["timestamp"],
            })
    return fps


def _empty_pauses(chunks, threshold_ms: float = 300) -> list[dict]:
    """CDM-2: Empty pauses (inter-word gaps >= threshold).

    Gaps between real words are measured by looking through [*] markers —
    the gap from the last real word's end to the next real word's start
    captures the full silence including any [*] span.
    """
    eps = []
    # Collect indices of real (non-[*]) words
    real_indices = [
        i for i, c in enumerate(chunks) if not _is_disfluency_marker(c["text"])
    ]
    for j in range(1, len(real_indices)):
        prev_i = real_indices[j - 1]
        cur_i = real_indices[j]
        prev_end = chunks[prev_i]["timestamp"]["end"]
        cur_start = chunks[cur_i]["timestamp"]["start"]
        gap_ms = (cur_start - prev_end) * 1000
        if gap_ms >= threshold_ms:
            eps.append({
                "after_index": prev_i,
                "gap_ms": round(gap_ms, 1),
                "start": prev_end,
                "end": cur_start,
            })
    return eps


def _empty_pauses_real(chunks, threshold_ms: float = 300) -> list[dict]:
    """Like _empty_pauses but for temporal calculations (same logic)."""
    return _empty_pauses(chunks, threshold_ms)


def _repetitions(chunks) -> list[dict]:
    """CDM-3: Consecutive identical words (1–3 word spans)."""
    words = [_clean_word(c["text"]) for c in chunks]
    reps = []
    i = 0
    while i < len(words):
        if not words[i] or _is_filler(chunks[i]["text"]):
            i += 1
            continue
        # check spans of length 1, 2, 3
        found = False
        for span in (3, 2, 1):
            if i + 2 * span > len(words):
                continue
            pattern = words[i:i + span]
            if all(pattern):
                candidate = words[i + span:i + 2 * span]
                if pattern == candidate:
                    reps.append({
                        "index": i,
                        "span": span,
                        "words": " ".join(pattern),
                    })
                    i += 2 * span
                    found = True
                    break
        if not found:
            i += 1
    return reps


def _retractions(chunks) -> list[dict]:
    """CDM-4: Consecutive same-category function words (heuristic)."""
    rts = []
    for i in range(1, len(chunks)):
        if _is_filler(chunks[i]["text"]) or _is_filler(chunks[i - 1]["text"]):
            continue
        cat_prev = _function_word_category(chunks[i - 1]["text"])
        cat_cur = _function_word_category(chunks[i]["text"])
        if cat_prev and cat_cur and cat_prev == cat_cur:
            w_prev = _clean_word(chunks[i - 1]["text"])
            w_cur = _clean_word(chunks[i]["text"])
            if w_prev != w_cur:  # same word = repetition, not retraction
                rts.append({
                    "index": i - 1,
                    "word1": w_prev,
                    "word2": w_cur,
                    "category": cat_prev,
                    "confidence": "medium",
                })
    return rts


def _temporal_markers(chunks, duration_s: float, clean_word_count: int):
    """Markers 5–9: speech rate, articulation rate, phonation ratio, pause stats, MLU."""
    # speaking segments (from first word start to last word end)
    if not chunks:
        return {}

    total_speaking_time = 0.0
    for c in chunks:
        # Exclude [*] disfluency markers from speaking time — they span
        # silence/hesitation gaps, not actual articulation
        if _is_disfluency_marker(c["text"]):
            continue
        ts = c["timestamp"]
        total_speaking_time += ts["end"] - ts["start"]

    # empty pauses for pause stats — use gaps between real words,
    # skipping over [*] markers to find true inter-word silences
    eps = _empty_pauses_real(chunks)
    pause_durations = [ep["gap_ms"] for ep in eps]

    # utterance segmentation: split at pauses >= 1000ms between real words
    utterances: list[list[str]] = []
    current_utt: list[str] = []
    prev_real_end: float | None = None
    for i, c in enumerate(chunks):
        if _is_disfluency_marker(c["text"]) or _is_filler(c["text"]):
            continue
        w = _clean_word(c["text"])
        if not w:
            continue
        # check if there's a long pause before this word (from last real word)
        if prev_real_end is not None:
            gap_ms = (c["timestamp"]["start"] - prev_real_end) * 1000
            if gap_ms >= 1000 and current_utt:
                utterances.append(current_utt)
                current_utt = []
        current_utt.append(w)
        prev_real_end = c["timestamp"]["end"]
    if current_utt:
        utterances.append(current_utt)

    # syllable count for syllable-based rates
    clean_words = _clean_tokens(chunks)
    total_syllables = sum(_syllable_count(w) for w in clean_words)

    speech_rate_wpm = (clean_word_count / duration_s) * 60 if duration_s > 0 else 0
    speech_rate_spm = (total_syllables / duration_s) * 60 if duration_s > 0 else 0
    art_rate_wpm = (
        (clean_word_count / total_speaking_time) * 60
        if total_speaking_time > 0 else 0
    )
    art_rate_spm = (
        (total_syllables / total_speaking_time) * 60
        if total_speaking_time > 0 else 0
    )
    phonation_ratio = (total_speaking_time / duration_s) * 100 if duration_s > 0 else 0

    utt_lengths = [len(u) for u in utterances]
    mlu = statistics.mean(utt_lengths) if utt_lengths else 0

    return {
        "speech_rate_wpm": round(speech_rate_wpm, 2),
        "speech_rate_spm": round(speech_rate_spm, 2),
        "articulation_rate_wpm": round(art_rate_wpm, 2),
        "articulation_rate_spm": round(art_rate_spm, 2),
        "phonation_time_ratio_pct": round(phonation_ratio, 2),
        "pause_mean_ms": round(statistics.mean(pause_durations), 1) if pause_durations else None,
        "pause_median_ms": round(statistics.median(pause_durations), 1) if pause_durations else None,
        "pause_std_ms": round(statistics.stdev(pause_durations), 1) if len(pause_durations) > 1 else None,
        "pause_count": len(pause_durations),
        "mlu_words": round(mlu, 2),
        "utterance_count": len(utterances),
        "total_syllables": total_syllables,
    }


def _lexical_diversity(chunks) -> dict:
    """Markers 10–11: TTR and MTLD."""
    tokens = _clean_tokens(chunks)
    if not tokens:
        return {"ttr": None, "mtld": None}
    types = set(tokens)
    ttr = len(types) / len(tokens)
    try:
        mtld = lex_div.mtld(tokens)
    except Exception:
        mtld = None
    return {"ttr": round(ttr, 4), "mtld": round(mtld, 2) if mtld is not None else None}


def _code_switching(chunks) -> list[dict]:
    """Marker 12: German words detected in Spanish speech."""
    switches = []
    for i, c in enumerate(chunks):
        w = _clean_word(c["text"])
        if w in GERMAN_WORDS:
            switches.append({"index": i, "word": w, **c["timestamp"]})
    return switches


def _disfluency_following(chunks, filled_pauses, empty_pauses) -> list[dict]:
    """Marker 13: Word class of the word following each filled/empty pause."""
    results = []
    for fp in filled_pauses:
        idx = fp["index"]
        if idx + 1 < len(chunks):
            next_word = chunks[idx + 1]["text"]
            results.append({
                "disfluency_type": "FP",
                "disfluency": fp["text"],
                "following_word": next_word.strip(),
                "following_class": _word_class(next_word),
            })
    for ep in empty_pauses:
        idx = ep["after_index"]
        if idx + 1 < len(chunks):
            next_word = chunks[idx + 1]["text"]
            results.append({
                "disfluency_type": "EP",
                "disfluency": f"[pause {ep['gap_ms']:.0f}ms]",
                "following_word": next_word.strip(),
                "following_class": _word_class(next_word),
            })
    return results


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------

def extract_markers(
    transcript: dict,
    duration_s: float | None = None,
    textgrid_path: str | None = None,
) -> dict:
    """Extract all markers from a transcript dict.

    Includes the original Schmid-style + temporal + lexical markers,
    Phase-1 acoustic markers (F0, voice quality, spectral tilt, MFCCs),
    and — when a TextGrid is provided — Phase-2 alignment-based markers
    (VOT, vowel formants, rhythm metrics).

    Parameters
    ----------
    transcript : dict
        Output of ``transcribe.transcribe()``.
    duration_s : float, optional
        Override audio duration (defaults to ``transcript["duration_s"]``).
    textgrid_path : str, optional
        Path to an MFA-produced TextGrid.  When provided, alignment-based
        markers (VOT, vowel formants/VSA, rhythm metrics) are extracted.

    Returns
    -------
    dict with keys for every marker group.
    """
    chunks = transcript["chunks"]
    if duration_s is None:
        duration_s = transcript["duration_s"]

    clean_words = _clean_tokens(chunks)
    word_count = len(clean_words)

    # A. Schmid disfluency markers
    fps = _filled_pauses(chunks)
    eps = _empty_pauses(chunks)
    rps = _repetitions(chunks)
    rts = _retractions(chunks)

    per100 = lambda n: round((n / word_count) * 100, 2) if word_count else 0

    # B. Temporal / fluency
    temporal = _temporal_markers(chunks, duration_s, word_count)

    # C. Lexical
    lex = _lexical_diversity(chunks)
    cs = _code_switching(chunks)

    # D. Disfluency context
    df_context = _disfluency_following(chunks, fps, eps)
    class_counts = Counter(d["following_class"] for d in df_context)

    # E. Acoustic markers (Phase 1)
    audio_path = transcript.get("audio_path")
    if audio_path:
        acoustic = extract_acoustic_markers(audio_path)
    else:
        acoustic = {}

    # F. Alignment-based markers (Phase 2) — requires MFA TextGrid
    if textgrid_path and audio_path:
        alignment = extract_alignment_markers(audio_path, textgrid_path)
    else:
        alignment = {}

    return {
        "word_count_clean": word_count,
        "duration_s": duration_s,
        # A — Schmid disfluency markers
        "filled_pauses": {
            "count": len(fps),
            "per_100_words": per100(len(fps)),
            "items": fps,
        },
        "empty_pauses": {
            "count": len(eps),
            "per_100_words": per100(len(eps)),
            "threshold_ms": 300,
            "items": eps,
        },
        "repetitions": {
            "count": len(rps),
            "per_100_words": per100(len(rps)),
            "items": rps,
        },
        "retractions": {
            "count": len(rts),
            "per_100_words": per100(len(rts)),
            "items": rts,
        },
        # B — Temporal / fluency
        "temporal": temporal,
        # C — Lexical
        "lexical_diversity": lex,
        "code_switching": {
            "count": len(cs),
            "items": cs,
        },
        # D — Disfluency context
        "disfluency_following": {
            "items": df_context,
            "class_distribution": dict(class_counts),
        },
        # E — Acoustic markers (Phase 1)
        "acoustic": acoustic,
        # F — Alignment-based markers (Phase 2)
        "alignment": alignment,
    }