alignment_markers.py

"""Phase 2 alignment-based markers: VOT, vowel formants, rhythm metrics.

Requires:
  - Montreal Forced Aligner (``mfa``) on PATH — install via conda:
        conda install -c conda-forge montreal-forced-aligner
        mfa model download acoustic spanish_mfa
        mfa model download dictionary spanish_mfa
  - praatio for TextGrid parsing (pip install praatio)
  - parselmouth (already installed in Phase 1)

Workflow:
  1. ``prepare_corpus``  — write .lab files from Whisper transcripts
  2. ``run_mfa_alignment`` — call ``mfa align`` via subprocess
  3. ``extract_alignment_markers`` — read TextGrid + audio → features
"""

from __future__ import annotations

import math
import shutil
import subprocess
import tempfile
from pathlib import Path

import numpy as np
import parselmouth
from parselmouth.praat import call
from praatio import textgrid as tgio

# ---------------------------------------------------------------------------
# Phone‐set definitions  (MFA Spanish IPA inventory)
# ---------------------------------------------------------------------------

VOWELS = {"a", "e", "i", "o", "u",
           "a\u02D0", "e\u02D0", "i\u02D0", "o\u02D0", "u\u02D0"}  # long variants

# Voiceless stops — primary VOT targets (Spanish short-lag → German long-lag)
# MFA Spanish uses dental t̪ and palatal c alongside plain p, k
VOICELESS_STOPS = {"p", "t", "k", "t\u032A", "c"}  # t̪ = dental t
# Voiced stops — secondary VOT targets (Spanish lead voicing)
# MFA Spanish uses dental d̪ alongside plain b, ɡ/g, and palatal ɟʝ
VOICED_STOPS = {"b", "d", "d\u032A", "\u0261", "g", "\u025F\u02DD"}  # d̪, ɡ, ɟʝ
# Voiced stop approximant allophones (excluded from VOT)
APPROXIMANTS = {"\u03B2", "\u00F0", "\u0263",       # β ð ɣ
                "\u0279"}                             # ɹ (if present)

# Corner vowels for Vowel Space Area
CORNER_VOWELS = {"a", "i", "u"}

# All consonant-like phones (anything not a vowel and not silence)
_SILENCE_LABELS = {"", "sil", "sp", "spn", "<eps>"}


def _is_vowel(phone: str) -> bool:
    """True if *phone* is a vowel (including long variants)."""
    return phone.lower().strip() in VOWELS


def _is_silence(phone: str) -> bool:
    return phone.lower().strip() in _SILENCE_LABELS


# ---------------------------------------------------------------------------
# 1.  Corpus preparation  (Whisper transcript → .lab files)
# ---------------------------------------------------------------------------

def prepare_corpus(
    transcripts: dict[str, dict],
    corpus_dir: str | Path,
    audio_dir: str | Path,
) -> Path:
    """Create an MFA-compatible corpus from Whisper transcripts.

    Parameters
    ----------
    transcripts : dict
        Mapping ``{speaker_id: transcript_dict}`` where each
        ``transcript_dict`` is the output of ``transcribe.transcribe()``.
    corpus_dir : path
        Directory to write the corpus into (created if needed).
    audio_dir : path
        Directory containing the original WAV files.

    Returns
    -------
    Path to the corpus directory.
    """
    corpus_dir = Path(corpus_dir)
    corpus_dir.mkdir(parents=True, exist_ok=True)
    audio_dir = Path(audio_dir)

    for speaker_id, transcript in transcripts.items():
        # Get clean text (strip disfluency markers, keep real words)
        words = []
        for chunk in transcript["chunks"]:
            text = chunk["text"].strip()
            if text and text != "[*]":
                # Strip brackets from CrisperWhisper tokens like [UH]
                clean = text.strip("[]")
                if clean:
                    words.append(clean)
        lab_text = " ".join(words)

        # Resolve audio file
        audio_path = Path(transcript["audio_path"])
        if not audio_path.is_absolute():
            audio_path = audio_dir / audio_path.name
        if not audio_path.exists():
            # Try matching by speaker ID
            candidates = list(audio_dir.glob(f"{speaker_id}*.[Ww][Aa][Vv]"))
            if candidates:
                audio_path = candidates[0]

        # Write symlink to audio + .lab file
        dest_wav = corpus_dir / f"{speaker_id}.wav"
        dest_lab = corpus_dir / f"{speaker_id}.lab"

        if not dest_wav.exists():
            # Symlink so we don't copy large files
            dest_wav.symlink_to(audio_path.resolve())

        dest_lab.write_text(lab_text, encoding="utf-8")

    return corpus_dir


# ---------------------------------------------------------------------------
# 2.  MFA alignment
# ---------------------------------------------------------------------------

def run_mfa_alignment(
    corpus_dir: str | Path,
    output_dir: str | Path,
    dictionary: str = "spanish_mfa",
    acoustic_model: str = "spanish_mfa",
    num_jobs: int = 4,
    clean: bool = True,
) -> Path:
    """Run ``mfa align`` on a prepared corpus.

    Returns the output directory containing TextGrid files.
    Raises RuntimeError if ``mfa`` is not found or alignment fails.
    """
    corpus_dir = Path(corpus_dir)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    mfa_bin = shutil.which("mfa")
    if mfa_bin is None:
        raise RuntimeError(
            "Montreal Forced Aligner (mfa) not found on PATH.\n"
            "Install via conda:\n"
            "  conda install -c conda-forge montreal-forced-aligner\n"
            "  mfa model download acoustic spanish_mfa\n"
            "  mfa model download dictionary spanish_mfa"
        )

    cmd = [
        mfa_bin, "align",
        str(corpus_dir),
        dictionary,
        acoustic_model,
        str(output_dir),
        "--output_format", "long_textgrid",
        "--num_jobs", str(num_jobs),
        "--single_speaker",
    ]
    if clean:
        cmd.append("--clean")

    result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
    if result.returncode != 0:
        raise RuntimeError(
            f"MFA alignment failed (exit {result.returncode}):\n"
            f"STDOUT:\n{result.stdout[-2000:]}\n"
            f"STDERR:\n{result.stderr[-2000:]}"
        )

    return output_dir


# ---------------------------------------------------------------------------
# 3.  TextGrid parsing
# ---------------------------------------------------------------------------

def parse_textgrid(tg_path: str | Path) -> dict:
    """Parse an MFA TextGrid into word and phone interval lists.

    Returns
    -------
    dict with keys ``"words"`` and ``"phones"``, each a list of
    ``(start, end, label)`` tuples.
    """
    tg = tgio.openTextgrid(str(tg_path), includeEmptyIntervals=False)

    phones = []
    words = []

    for tier_name in tg.tierNames:
        tier = tg.getTier(tier_name)
        lower = tier_name.lower()
        entries = [(float(s), float(e), label)
                   for s, e, label in tier.entries]
        if "phone" in lower:
            phones = entries
        elif "word" in lower:
            words = entries

    return {"words": words, "phones": phones}


# ---------------------------------------------------------------------------
# 4a.  VOT extraction
# ---------------------------------------------------------------------------

def _detect_burst_time(
    sound: parselmouth.Sound,
    start: float,
    end: float,
) -> float | None:
    """Detect the burst release within a stop-consonant interval.

    Uses the intensity contour: the burst is the sharpest intensity
    rise within the stop interval.  Returns the absolute time of the
    burst, or None if detection fails.
    """
    duration = end - start
    if duration < 0.010:
        return None

    try:
        segment = sound.extract_part(start, end,
                                     parselmouth.WindowShape.RECTANGULAR,
                                     1.0, False)
        intensity = segment.to_intensity(minimum_pitch=400, time_step=0.0005)
    except Exception:
        return None

    n = intensity.get_number_of_frames()
    if n < 3:
        return None

    # Find the frame with the largest intensity rise (derivative)
    best_rise = -np.inf
    best_time = None
    for i in range(2, n + 1):
        t_prev = intensity.get_time_from_frame_number(i - 1)
        t_cur = intensity.get_time_from_frame_number(i)
        val_prev = intensity.get_value(t_prev)
        val_cur = intensity.get_value(t_cur)

        if np.isnan(val_prev) or np.isnan(val_cur):
            continue

        rise = val_cur - val_prev
        if rise > best_rise:
            best_rise = rise
            best_time = t_cur

    if best_time is None or best_rise < 1.0:  # minimum 1 dB rise
        return None

    # Convert back to absolute time
    return start + best_time


def extract_vot(
    phones: list[tuple[float, float, str]],
    sound: parselmouth.Sound,
) -> dict:
    """Measure Voice Onset Time for voiceless and voiced stops.

    MFA places stop and vowel boundaries contiguously (no gap), so the
    stop interval itself contains both closure and release/aspiration.

    For voiceless stops /p, t, k/:
      VOT = stop_end − burst_time  (positive: time from burst to vowel onset).
    For voiced stops /b, d, g/:
      VOT = stop_end − burst_time  (may be short if voicing starts early).

    When burst detection fails, VOT is estimated as a fraction of the
    stop duration (the release portion, typically the final ~40%).

    Returns per-phone-type mean VOT and individual measurements.
    """
    measurements: list[dict] = []

    for i, (start, end, label) in enumerate(phones):
        phone = label.lower().strip()

        # Only measure stops followed by a vowel
        if phone not in VOICELESS_STOPS and phone not in VOICED_STOPS:
            continue

        # Find the next non-silence phone
        next_phone = None
        for j in range(i + 1, len(phones)):
            _, _, nlabel = phones[j]
            if not _is_silence(nlabel):
                next_phone = phones[j]
                break

        if next_phone is None or not _is_vowel(next_phone[2]):
            continue

        vowel_start = next_phone[0]
        stop_type = "voiceless" if phone in VOICELESS_STOPS else "voiced"
        stop_dur_ms = (end - start) * 1000

        # Detect burst within the stop interval
        burst_time = _detect_burst_time(sound, start, end)

        if burst_time is not None:
            # VOT = time from burst to stop/vowel boundary
            vot_ms = (end - burst_time) * 1000
        else:
            # Fallback: estimate VOT as the final 40% of the stop
            # (closure is ~60%, release+aspiration is ~40%)
            vot_ms = stop_dur_ms * 0.4
            burst_time = start + (end - start) * 0.6

        # Sanity: VOT shouldn't exceed the stop duration
        vot_ms = min(vot_ms, stop_dur_ms)
        vot_ms = max(vot_ms, 0.0)

        measurements.append({
            "phone": phone,
            "type": stop_type,
            "position_s": round(start, 3),
            "stop_dur_ms": round(stop_dur_ms, 2),
            "vot_ms": round(vot_ms, 2),
            "burst_time_s": round(burst_time, 4),
            "vowel_onset_s": round(vowel_start, 4),
        })

    # Aggregate by phone and type
    voiceless_vots = [m["vot_ms"] for m in measurements
                      if m["type"] == "voiceless"]
    voiced_vots = [m["vot_ms"] for m in measurements
                   if m["type"] == "voiced"]

    per_phone: dict[str, list[float]] = {}
    for m in measurements:
        per_phone.setdefault(m["phone"], []).append(m["vot_ms"])

    phone_means = {p: round(float(np.mean(vs)), 2)
                   for p, vs in per_phone.items()}

    return {
        "voiceless_mean_vot_ms": round(float(np.mean(voiceless_vots)), 2)
            if voiceless_vots else None,
        "voiceless_sd_vot_ms": round(float(np.std(voiceless_vots, ddof=1)), 2)
            if len(voiceless_vots) > 1 else None,
        "voiced_mean_vot_ms": round(float(np.mean(voiced_vots)), 2)
            if voiced_vots else None,
        "voiced_sd_vot_ms": round(float(np.std(voiced_vots, ddof=1)), 2)
            if len(voiced_vots) > 1 else None,
        "per_phone_mean_ms": phone_means,
        "n_voiceless": len(voiceless_vots),
        "n_voiced": len(voiced_vots),
        "measurements": measurements,
    }


# ---------------------------------------------------------------------------
# 4b.  Vowel formant extraction + Vowel Space Area
# ---------------------------------------------------------------------------

def extract_vowel_formants(
    phones: list[tuple[float, float, str]],
    sound: parselmouth.Sound,
    max_formant: float = 5500.0,
    n_formants: int = 5,
) -> dict:
    """Extract F1/F2/F3 at the temporal midpoint of each vowel.

    Parameters
    ----------
    max_formant : float
        Maximum formant frequency for Burg analysis.
        Use 5500 for female speakers, 5000 for male speakers.
        Default 5500 (conservative for mixed/unknown gender).

    Returns per-vowel-type mean formants, all individual measurements,
    and Vowel Space Area computed from corner vowels /a, i, u/.
    """
    formant_obj = call(sound, "To Formant (burg)",
                       0.025, n_formants, max_formant, 0.025, 50.0)

    measurements: list[dict] = []

    for start, end, label in phones:
        phone = label.lower().strip()
        if not _is_vowel(phone):
            continue

        # Strip length marks for grouping
        vowel_id = phone.replace("\u02D0", "")
        duration = end - start
        if duration < 0.02:
            continue

        midpoint = (start + end) / 2

        f1 = call(formant_obj, "Get value at time", 1, midpoint, "Hertz", "Linear")
        f2 = call(formant_obj, "Get value at time", 2, midpoint, "Hertz", "Linear")
        f3 = call(formant_obj, "Get value at time", 3, midpoint, "Hertz", "Linear")

        if np.isnan(f1) or np.isnan(f2):
            continue

        measurements.append({
            "vowel": vowel_id,
            "midpoint_s": round(midpoint, 4),
            "duration_ms": round(duration * 1000, 1),
            "f1_hz": round(f1, 1),
            "f2_hz": round(f2, 1),
            "f3_hz": round(f3, 1) if not np.isnan(f3) else None,
        })

    # Per-vowel means
    vowel_data: dict[str, list[dict]] = {}
    for m in measurements:
        vowel_data.setdefault(m["vowel"], []).append(m)

    per_vowel: dict[str, dict] = {}
    for v, items in vowel_data.items():
        f1s = [it["f1_hz"] for it in items]
        f2s = [it["f2_hz"] for it in items]
        per_vowel[v] = {
            "n": len(items),
            "f1_mean_hz": round(float(np.mean(f1s)), 1),
            "f1_sd_hz": round(float(np.std(f1s, ddof=1)), 1) if len(f1s) > 1 else None,
            "f2_mean_hz": round(float(np.mean(f2s)), 1),
            "f2_sd_hz": round(float(np.std(f2s, ddof=1)), 1) if len(f2s) > 1 else None,
        }

    # Vowel Space Area (triangle: /a/, /i/, /u/)
    vsa = _compute_vsa(per_vowel)

    # Vowel Formant Dispersion
    vfd = _compute_vfd(per_vowel)

    return {
        "per_vowel": per_vowel,
        "vowel_space_area": vsa,
        "vowel_formant_dispersion": vfd,
        "n_total": len(measurements),
        "measurements": measurements,
    }


def _compute_vsa(per_vowel: dict[str, dict]) -> float | None:
    """Vowel Space Area — triangle formed by /a/, /i/, /u/ in F1×F2 space.

    Uses the Shoelace formula for the area of a triangle:
      VSA = 0.5 * |F1a(F2i - F2u) + F1i(F2u - F2a) + F1u(F2a - F2i)|
    """
    corners = {}
    for v in CORNER_VOWELS:
        if v in per_vowel:
            corners[v] = (per_vowel[v]["f1_mean_hz"],
                          per_vowel[v]["f2_mean_hz"])

    if len(corners) < 3:
        return None

    a = corners["a"]
    i = corners["i"]
    u = corners["u"]

    area = 0.5 * abs(
        a[0] * (i[1] - u[1]) +
        i[0] * (u[1] - a[1]) +
        u[0] * (a[1] - i[1])
    )
    return round(area, 1)


def _compute_vfd(per_vowel: dict[str, dict]) -> float | None:
    """Vowel Formant Dispersion — mean Euclidean distance from centroid."""
    if not per_vowel:
        return None

    f1_all = [v["f1_mean_hz"] for v in per_vowel.values()]
    f2_all = [v["f2_mean_hz"] for v in per_vowel.values()]

    centroid_f1 = np.mean(f1_all)
    centroid_f2 = np.mean(f2_all)

    distances = []
    for v in per_vowel.values():
        d = math.sqrt((v["f1_mean_hz"] - centroid_f1) ** 2 +
                      (v["f2_mean_hz"] - centroid_f2) ** 2)
        distances.append(d)

    return round(float(np.mean(distances)), 1)


# ---------------------------------------------------------------------------
# 4c.  Rhythm metrics
# ---------------------------------------------------------------------------

def extract_rhythm_metrics(
    phones: list[tuple[float, float, str]],
) -> dict:
    """Compute rhythm metrics from phone-level intervals.

    Returns %V, deltaC, deltaV, VarcoC, VarcoV, rPVI-C, nPVI-V.
    """
    # Step 1: classify each phone as C or V, skip silence
    cv_intervals: list[tuple[float, float, str]] = []
    for start, end, label in phones:
        if _is_silence(label):
            continue
        category = "V" if _is_vowel(label) else "C"
        cv_intervals.append((start, end, category))

    if not cv_intervals:
        return _empty_rhythm()

    # Step 2: merge adjacent same-type intervals
    merged: list[tuple[float, float, str]] = [cv_intervals[0]]
    for start, end, cat in cv_intervals[1:]:
        prev_start, prev_end, prev_cat = merged[-1]
        if cat == prev_cat and abs(start - prev_end) < 0.001:
            # Merge
            merged[-1] = (prev_start, end, cat)
        else:
            merged.append((start, end, cat))

    # Step 3: compute durations
    c_durations = [(e - s) * 1000 for s, e, cat in merged if cat == "C"]
    v_durations = [(e - s) * 1000 for s, e, cat in merged if cat == "V"]

    if len(c_durations) < 2 or len(v_durations) < 2:
        return _empty_rhythm()

    c_arr = np.array(c_durations)
    v_arr = np.array(v_durations)
    total_dur = sum(c_durations) + sum(v_durations)

    # %V  — proportion of vocalic intervals
    pct_v = (sum(v_durations) / total_dur) * 100 if total_dur > 0 else 0

    # deltaC, deltaV  — standard deviations
    delta_c = float(np.std(c_arr, ddof=1))
    delta_v = float(np.std(v_arr, ddof=1))

    # VarcoC, VarcoV  — variation coefficients (rate-normalized)
    mean_c = float(np.mean(c_arr))
    mean_v = float(np.mean(v_arr))
    varco_c = (delta_c / mean_c) * 100 if mean_c > 0 else 0
    varco_v = (delta_v / mean_v) * 100 if mean_v > 0 else 0

    # rPVI-C  — raw Pairwise Variability Index for consonants
    rpvi_c = float(np.mean(np.abs(np.diff(c_arr))))

    # nPVI-V  — normalized PVI for vowels
    npvi_v = _npvi(v_arr)

    return {
        "pct_v": round(pct_v, 2),
        "delta_c_ms": round(delta_c, 2),
        "delta_v_ms": round(delta_v, 2),
        "varco_c": round(varco_c, 2),
        "varco_v": round(varco_v, 2),
        "rpvi_c": round(rpvi_c, 2),
        "npvi_v": round(npvi_v, 2),
        "n_c_intervals": len(c_durations),
        "n_v_intervals": len(v_durations),
        "mean_c_ms": round(mean_c, 2),
        "mean_v_ms": round(mean_v, 2),
    }


def _npvi(durations: np.ndarray) -> float:
    """Normalized Pairwise Variability Index.

    nPVI = 100 * (1/(n-1)) * Σ |d_k - d_{k+1}| / ((d_k + d_{k+1}) / 2)
    """
    n = len(durations)
    if n < 2:
        return 0.0
    total = 0.0
    for k in range(n - 1):
        avg = (durations[k] + durations[k + 1]) / 2
        if avg > 0:
            total += abs(durations[k] - durations[k + 1]) / avg
    return 100.0 * total / (n - 1)


def _empty_rhythm() -> dict:
    return {
        "pct_v": None, "delta_c_ms": None, "delta_v_ms": None,
        "varco_c": None, "varco_v": None,
        "rpvi_c": None, "npvi_v": None,
        "n_c_intervals": 0, "n_v_intervals": 0,
        "mean_c_ms": None, "mean_v_ms": None,
    }


# ---------------------------------------------------------------------------
# 5.  Public entry point
# ---------------------------------------------------------------------------

def extract_alignment_markers(
    audio_path: str | Path,
    textgrid_path: str | Path,
    max_formant: float = 5500.0,
) -> dict:
    """Extract all Phase-2 alignment-based markers.

    Parameters
    ----------
    audio_path : path
        Path to the WAV file.
    textgrid_path : path
        Path to the MFA-produced TextGrid.
    max_formant : float
        Maximum formant frequency for Burg analysis (5500 for female,
        5000 for male, default 5500).

    Returns
    -------
    dict with keys ``"vot"``, ``"vowel_formants"``, ``"rhythm"``.
    """
    sound = parselmouth.Sound(str(audio_path))
    tg_data = parse_textgrid(textgrid_path)
    phones = tg_data["phones"]

    if not phones:
        return {"vot": {}, "vowel_formants": {}, "rhythm": _empty_rhythm()}

    vot = extract_vot(phones, sound)
    formants = extract_vowel_formants(phones, sound, max_formant=max_formant)
    rhythm = extract_rhythm_metrics(phones)

    return {
        "vot": vot,
        "vowel_formants": formants,
        "rhythm": rhythm,
    }