core.py

"""
Phonemlyze core pipeline:
  1. Whisper  → spoken text + transcription confidence
  2. eSpeak-ng → target IPA  (theoretically correct pronunciation)
  3. Wav2Vec2  → actual IPA  (physically spoken phonemes) + per-phoneme confidence
  4. Alignment → phoneme-level diff with confidence scores
"""

import subprocess
import os
import re
import wave
import tempfile
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import librosa
import torch
import whisper
from transformers import (
    Wav2Vec2FeatureExtractor,
    Wav2Vec2PhonemeCTCTokenizer,
    Wav2Vec2Processor,
    Wav2Vec2ForCTC,
)
import difflib


# ── Types ─────────────────────────────────────────────────────────────────────

Status = str  # "match" | "replace" | "delete" | "insert"

@dataclass
class AlignedPair:
    target: str
    actual: str
    status: Status
    confidence: Optional[float] = None

@dataclass
class AnalysisResult:
    text: str
    transcription_confidence: float
    target_ipa: str
    actual_ipa: str
    model_confidence: float
    alignment: List[AlignedPair]
    match_rate: float
    errors: List[AlignedPair] = field(default_factory=list)


# ── Main class ────────────────────────────────────────────────────────────────

class Phonemlyze:
    WAV2VEC2_MODEL = "facebook/wav2vec2-lv-60-espeak-cv-ft"

    def __init__(self, whisper_model: str = "base", device: Optional[str] = None):
        if device:
            self.device = device
        elif torch.cuda.is_available():
            self.device = "cuda"
        elif torch.backends.mps.is_available():
            self.device = "mps"
        else:
            self.device = "cpu"

        print(f"[Phonemlyze] Device: {self.device}")
        print(f"[Phonemlyze] Loading Whisper ({whisper_model}) …")
        whisper_device = "cpu" if self.device == "mps" else self.device
        self.whisper = whisper.load_model(whisper_model, device=whisper_device)

        print(f"[Phonemlyze] Loading Wav2Vec2 ({self.WAV2VEC2_MODEL}) …")
        feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(self.WAV2VEC2_MODEL)
        # do_phonemize=False skips the espeak backend init — we only need the
        # tokenizer for decoding output IDs, not for phonemizing input text
        tokenizer = Wav2Vec2PhonemeCTCTokenizer.from_pretrained(
            self.WAV2VEC2_MODEL, do_phonemize=False
        )
        self.processor = Wav2Vec2Processor(
            feature_extractor=feature_extractor, tokenizer=tokenizer
        )
        self.wav2vec2 = Wav2Vec2ForCTC.from_pretrained(self.WAV2VEC2_MODEL).to(self.device)
        self.wav2vec2.eval()
        self.blank_id = self.wav2vec2.config.pad_token_id

        # Build sorted phoneme list (longest first) for IPA segmentation
        _tok = getattr(self.processor, "tokenizer", None)
        if _tok is not None:
            _special = set(getattr(_tok, "all_special_tokens", []))
            self._phoneme_vocab: List[str] = sorted(
                {t for t in _tok.get_vocab() if t not in _special and t not in ("|", "<pad>")},
                key=len,
                reverse=True,
            )
        else:
            self._phoneme_vocab = []

        print("[Phonemlyze] Ready.")

    # ── Stage 1: speech → text + confidence ───────────────────────────────

    def transcribe(self, audio_path: str) -> Tuple[str, float]:
        result = self.whisper.transcribe(audio_path, fp16=(self.device == "cuda"))
        text = result["text"].strip()
        segments = result.get("segments", [])
        if segments:
            avg_logprob = float(np.mean([s["avg_logprob"] for s in segments]))
            confidence = float(np.clip(np.exp(avg_logprob), 0.0, 1.0))
        else:
            confidence = 0.0
        return text, confidence

    # ── Stage 2: text → target IPA via eSpeak-ng ──────────────────────────

    _ESPEAK_BINARIES = [
        "/opt/homebrew/bin/espeak-ng",
        "/usr/local/bin/espeak-ng",
        "espeak-ng",
    ]

    def _run_espeak(self, args: List[str]) -> Optional[subprocess.CompletedProcess]:
        """Run espeak-ng with the given args, trying known binary locations."""
        for binary in self._ESPEAK_BINARIES:
            if binary == "espeak-ng" or os.path.isfile(binary):
                try:
                    return subprocess.run([binary] + args, capture_output=True, text=True)
                except FileNotFoundError:
                    continue
        return None

    # Map our lang codes to Piper voice names (rhasspy/piper-voices on HF Hub)
    # Languages without an entry fall back to eSpeak-ng synthesis automatically.
    PIPER_VOICES: Dict[str, str] = {
        "ar":    "ar_JO-kareem-medium",       # Jordanian Arabic (only Arabic voice available)
        "cs":    "cs_CZ-jirka-medium",
        "de":    "de_DE-thorsten-medium",
        "el":    "el_GR-rapunzelina-medium",
        "en-gb": "en_GB-alan-medium",         # Received Pronunciation (BBC standard)
        "en-us": "en_US-lessac-medium",       # General American (broadcast standard)
        "es-419":"es_MX-ald-medium",          # Mexican Spanish (neutral Latin American broadcast)
        "es-es": "es_ES-sharvard-medium",     # Castilian
        "fr":    "fr_FR-tom-medium",          # Parisian French
        # fr-ca: no Piper voice available — falls back to eSpeak
        "hi":    "hi_IN-rohan-medium",
        "hu":    "hu_HU-anna-medium",
        "it":    "it_IT-riccardo-x_low",
        # ja: no Piper voice available — falls back to eSpeak
        "nl":    "nl_NL-mls-medium",
        "pt-br": "pt_BR-faber-medium",
        "ru":    "ru_RU-dmitri-medium",
        # th: no Piper voice available — falls back to eSpeak
        "tr":    "tr_TR-dfki-medium",
        "uk":    "uk_UA-ukrainian_tts-medium",
        "zh":    "zh_CN-huayan-medium",
        "zh-tw": "zh_CN-huayan-medium",       # No TW-specific voice; Mandarin phonology is identical
        # uk: Ukrainian (above)
    }

    def _load_piper_voice(self, lang: str) -> Any:
        """Lazily download (once) and cache a Piper voice for the given language."""
        if not hasattr(self, "_piper_cache"):
            self._piper_cache: Dict[str, Any] = {}
        if lang in self._piper_cache:
            return self._piper_cache[lang]

        voice_name = self.PIPER_VOICES.get(lang) or self.PIPER_VOICES["en-us"]
        locale, speaker, quality = voice_name.split("-", 2)
        lang_prefix = locale.split("_")[0]
        subfolder = f"{lang_prefix}/{locale}/{speaker}/{quality}"

        try:
            from piper.voice import PiperVoice
            from huggingface_hub import hf_hub_download

            print(f"[Phonemlyze] Downloading Piper voice '{voice_name}' …", flush=True)
            onnx_path = hf_hub_download(
                repo_id="rhasspy/piper-voices",
                filename=f"{subfolder}/{voice_name}.onnx",
            )
            config_path = hf_hub_download(
                repo_id="rhasspy/piper-voices",
                filename=f"{subfolder}/{voice_name}.onnx.json",
            )
            voice = PiperVoice.load(onnx_path, config_path=config_path, use_cuda=False)
            print(f"[Phonemlyze] Piper voice '{voice_name}' ready.", flush=True)
            self._piper_cache[lang] = voice
        except Exception as exc:
            print(f"[Phonemlyze] Piper TTS unavailable ({exc}), falling back to eSpeak.", flush=True)
            self._piper_cache[lang] = None

        return self._piper_cache[lang]

    def synthesize_ipa(self, ipa: str, lang: str = "en-us") -> Optional[str]:
        """Synthesize a space-separated IPA phoneme string to audio.

        Uses Piper's phoneme_ids_to_audio to feed phonemes directly into the
        neural vocoder, bypassing any text-to-phoneme step (which would cause
        eSpeak to read each symbol as a letter name).  Falls back to eSpeak for
        languages without a Piper voice.
        """
        if not ipa.strip():
            return None

        phonemes = self._tokenize(ipa)   # handles word-grouped format correctly
        fd, out_path = tempfile.mkstemp(suffix=".wav")
        os.close(fd)

        # ── Piper path: phonemes → IDs → audio ────────────────────────────
        voice = self._load_piper_voice(lang)
        if voice is not None:
            try:
                # Piper's vocab is single characters only — split multi-char
                # tokens (dʒ, eɪ, oʊ …) into their individual characters,
                # keeping only characters that exist in the model's vocab.
                vocab = set(voice.config.phoneme_id_map.keys())
                expanded: List[str] = []
                for ph in phonemes:
                    if ph in vocab:
                        expanded.append(ph)
                    else:
                        expanded.extend(c for c in ph if c in vocab)

                phoneme_ids = voice.phonemes_to_ids(expanded)
                audio = voice.phoneme_ids_to_audio(phoneme_ids)
                # Normalise float32 → int16 WAV
                max_val = np.max(np.abs(audio))
                if max_val > 0:
                    audio = audio / max_val
                audio_int16 = (np.clip(audio, -1.0, 1.0) * 32767).astype(np.int16)
                with wave.open(out_path, "wb") as wf:
                    wf.setnchannels(1)
                    wf.setsampwidth(2)
                    wf.setframerate(voice.config.sample_rate)
                    wf.writeframes(audio_int16.tobytes())
                return out_path
            except Exception as exc:
                print(f"[Phonemlyze] Piper IPA synthesis failed ({exc}), falling back to eSpeak", flush=True)

        # ── eSpeak fallback ────────────────────────────────────────────────
        # Join tokens without spaces and hope eSpeak reads the IPA characters
        result = self._run_espeak(["-v", lang, "-w", out_path, "".join(phonemes)])
        return out_path if result is not None else None

    def synthesize(self, text: str, lang: str = "en") -> Optional[str]:
        """Synthesize text to a WAV file using Piper TTS (falls back to eSpeak). Returns path."""
        fd, out_path = tempfile.mkstemp(suffix=".wav")
        os.close(fd)

        voice = self._load_piper_voice(lang)
        if voice is not None:
            with wave.open(out_path, "wb") as wav_file:
                voice.synthesize_wav(text, wav_file)
            return out_path

        # Fallback: eSpeak-ng
        result = self._run_espeak(["-v", lang, "-w", out_path, text])
        return out_path if result is not None else None

    def text_to_ipa(self, text: str, lang: str = "en") -> str:
        proc = self._run_espeak(["-q", "--ipa", "-v", lang, text])
        if proc is None:
            raise RuntimeError("espeak-ng not found at any known location")
        raw = (proc.stdout or proc.stderr).strip()
        return self._segment_ipa(self._clean_espeak_ipa(raw))

    @staticmethod
    def _normalise_ipa(text: str) -> str:
        """Shared normalisation applied to both eSpeak and Wav2Vec2 output.

        eSpeak uses '_' between phonemes *within* a word and space between
        words.  Removing '_' (not replacing with space) preserves that
        word-grouping, giving e.g. 'kəmɪl nændʒɑni' instead of
        'k ə m ɪ l n æ n dʒ ɑ n i'.
        """
        text = text.replace("_", "")        # join phonemes within eSpeak words
        text = text.replace("|", " ")        # word-boundary → space
        text = re.sub(r"\([a-z-]+\)", "", text)  # strip lang-switch tags
        text = re.sub(r"[ˈˌːˑ]", "", text)  # strip stress / length marks
        text = re.sub(r"\s+", " ", text).strip()
        return text

    @staticmethod
    def _clean_espeak_ipa(raw: str) -> str:
        lines = [line.strip() for line in raw.splitlines() if line.strip()]
        return Phonemlyze._normalise_ipa(" ".join(lines))

    def _segment_word(self, word: str) -> List[str]:
        """Greedily split one IPA word-string into individual phoneme tokens."""
        if not self._phoneme_vocab:
            return list(word)
        result: List[str] = []
        i = 0
        while i < len(word):
            for token in self._phoneme_vocab:
                if word[i:].startswith(token):
                    result.append(token)
                    i += len(token)
                    break
            else:
                result.append(word[i])
                i += 1
        return result

    def _inject_word_boundaries(self, actual_ipa: str, target_ipa: str) -> str:
        """Re-inject approximate word boundaries into a flat actual_ipa string.

        When Wav2Vec2 does not emit '|' tokens the whole utterance arrives as
        one unseparated chunk.  We use the target word phoneme-counts as a
        proportional guide to split the actual token sequence.

        Only activates when actual_ipa contains no spaces but target_ipa does.
        """
        if " " in actual_ipa:
            return actual_ipa          # already has boundaries — nothing to do
        target_words = target_ipa.split()
        if len(target_words) <= 1:
            return actual_ipa          # single-word target — nothing to split

        actual_tokens = self._tokenize(actual_ipa)
        if not actual_tokens:
            return actual_ipa

        target_counts = [len(self._segment_word(w)) for w in target_words]
        total_target  = sum(target_counts) or 1
        total_actual  = len(actual_tokens)
        n_words       = len(target_words)

        result_words: List[str] = []
        pos = 0
        for i, count in enumerate(target_counts):
            if i == n_words - 1:
                # Last word gets everything remaining
                result_words.append("".join(actual_tokens[pos:]))
            else:
                # Proportional share, but guarantee ≥1 phoneme per word
                n = max(1, round(total_actual * count / total_target))
                remaining_words = n_words - i - 1
                n = min(n, total_actual - pos - remaining_words)
                result_words.append("".join(actual_tokens[pos : pos + n]))
                pos += n

        return " ".join(w for w in result_words if w)

    def _segment_ipa(self, ipa: str) -> str:
        """Segment eSpeak IPA, preserving word groupings.

        Phonemes within a word are joined (no separator); words are separated
        by a single space.  E.g. 'kəmɪl nændʒɑni' — copyable and readable.
        For alignment use _tokenize() which further splits each word.
        """
        words_out = ["".join(self._segment_word(w)) for w in ipa.split()]
        return " ".join(words_out)

    # ── Stage 3: audio → actual IPA + per-phoneme confidence ──────────────

    def audio_to_ipa(self, audio_path: str) -> Tuple[str, Dict[int, float], float]:
        """Returns (ipa_string, {phoneme_index: confidence}, avg_confidence)."""
        speech, _ = librosa.load(audio_path, sr=16_000, mono=True)

        inputs = self.processor(
            speech,
            sampling_rate=16_000,
            return_tensors="pt",
            padding=True,
        )

        input_values = inputs.input_values.to(self.device)
        model_kwargs = {}
        attention_mask = inputs.get("attention_mask")
        if attention_mask is not None:
            model_kwargs["attention_mask"] = attention_mask.to(self.device)

        with torch.inference_mode():
            logits = self.wav2vec2(input_values, **model_kwargs).logits
            probs = torch.softmax(logits, dim=-1)
            predicted_ids = torch.argmax(logits, dim=-1)

        ids = predicted_ids[0].detach().cpu().tolist()
        frame_conf = probs[0].max(dim=-1).values.detach().cpu().tolist()

        tokenizer = getattr(self.processor, "tokenizer", None)
        special_tokens = set(getattr(tokenizer, "all_special_tokens", [])) if tokenizer else set()

        # CTC collapse: group consecutive same non-blank IDs
        grouped: List[Tuple[int, List[float]]] = []
        current_id: Optional[int] = None
        current_confs: List[float] = []

        for token_id, conf in zip(ids, frame_conf):
            token_id = int(token_id)
            if token_id == self.blank_id:
                if current_id is not None:
                    grouped.append((current_id, current_confs))
                    current_id = None
                    current_confs = []
                continue
            if token_id == current_id:
                current_confs.append(float(conf))
            else:
                if current_id is not None:
                    grouped.append((current_id, current_confs))
                current_id = token_id
                current_confs = [float(conf)]

        if current_id is not None:
            grouped.append((current_id, current_confs))

        # Convert IDs → phoneme strings; use '|' as word boundary
        words: List[List[str]] = [[]]
        word_confs: List[List[float]] = [[]]

        for token_id, confs in grouped:
            token = (
                tokenizer.convert_ids_to_tokens(token_id)
                if tokenizer is not None
                else self.wav2vec2.config.id2label.get(token_id, "")
            )
            if not token or token in special_tokens:
                continue
            if token == "|":
                if words[-1]:          # start a new word only if current one is non-empty
                    words.append([])
                    word_confs.append([])
            else:
                words[-1].append(token)
                word_confs[-1].append(float(np.mean(confs)))

        # Drop empty trailing word slots
        non_empty = [(w, c) for w, c in zip(words, word_confs) if w]
        if non_empty:
            words, word_confs = zip(*non_empty)  # type: ignore[assignment]
        else:
            words, word_confs = [], []

        # Word-grouped string: phonemes joined within words, space between words
        actual_ipa = " ".join("".join(w) for w in words)
        actual_ipa = self._normalise_ipa(actual_ipa)   # strips stress marks etc.

        # Flat list for confidence indexing
        phoneme_confs: List[float] = [c for wc in word_confs for c in wc]
        confidences = {i: round(c, 4) for i, c in enumerate(phoneme_confs)}
        avg_conf = float(np.mean(phoneme_confs)) if phoneme_confs else 0.0
        return actual_ipa, confidences, avg_conf

    # ── Stage 4: phoneme alignment with confidence ─────────────────────────

    def _tokenize(self, ipa: str) -> List[str]:
        """Split a word-grouped IPA string into individual phoneme tokens.

        Each space-separated chunk is re-segmented using the vocab so that
        multi-char phonemes like 'dʒ' or 'eɪ' are kept as single units.
        """
        tokens: List[str] = []
        for word in ipa.split():
            tokens.extend(self._segment_word(word))
        return tokens

    def align(
        self,
        target_ipa: str,
        actual_ipa: str,
        confidences: Dict[int, float],
    ) -> Tuple[List[AlignedPair], float]:
        target_phones = self._tokenize(target_ipa)
        actual_phones = self._tokenize(actual_ipa)

        matcher = difflib.SequenceMatcher(None, target_phones, actual_phones, autojunk=False)
        pairs: List[AlignedPair] = []
        matches = 0
        total = 0

        for op, i1, i2, j1, j2 in matcher.get_opcodes():
            t_chunk = target_phones[i1:i2]
            a_chunk = actual_phones[j1:j2]

            if op == "equal":
                for t, a, aj in zip(t_chunk, a_chunk, range(j1, j2)):
                    pairs.append(AlignedPair(t, a, "match", confidences.get(aj)))
                    matches += 1
                    total += 1

            elif op == "replace":
                n = max(len(t_chunk), len(a_chunk))
                for i in range(n):
                    t = t_chunk[i] if i < len(t_chunk) else "∅"
                    a = a_chunk[i] if i < len(a_chunk) else "∅"
                    pairs.append(AlignedPair(t, a, "replace", confidences.get(j1 + i)))
                    total += 1

            elif op == "delete":
                for t in t_chunk:
                    pairs.append(AlignedPair(t, "∅", "delete", None))
                    total += 1

            elif op == "insert":
                for i, a in enumerate(a_chunk):
                    pairs.append(AlignedPair("∅", a, "insert", confidences.get(j1 + i)))
                    total += 1

        match_rate = matches / total if total else 0.0
        return pairs, match_rate

    # ── Full pipeline ──────────────────────────────────────────────────────

    def analyze(
        self,
        audio_path: Optional[str],
        lang: str = "en",
        text_override: Optional[str] = None,
    ) -> AnalysisResult:
        # Use supplied text, or transcribe audio, whichever is available
        if text_override and text_override.strip():
            text = text_override.strip()
            transcription_confidence = 1.0  # user-supplied, fully trusted
        elif audio_path:
            text, transcription_confidence = self.transcribe(audio_path)
        else:
            raise ValueError("Provide either audio or text input.")

        target_ipa = self.text_to_ipa(text, lang)

        if audio_path:
            actual_ipa, confidences, model_confidence = self.audio_to_ipa(audio_path)
            # Re-inject word boundaries when Wav2Vec2 didn't emit '|' tokens
            if actual_ipa and target_ipa:
                actual_ipa = self._inject_word_boundaries(actual_ipa, target_ipa)
        else:
            actual_ipa, confidences, model_confidence = "", {}, 0.0

        with open("/tmp/phonemlyze_debug.txt", "w") as f:
            f.write(f"text:       {text!r}\n")
            f.write(f"audio_path: {audio_path!r}\n")
            f.write(f"target_ipa: {target_ipa!r}\n")
            f.write(f"target_tok: {self._tokenize(target_ipa)}\n")
            f.write(f"actual_ipa: {actual_ipa!r}\n")
            f.write(f"actual_tok: {self._tokenize(actual_ipa)}\n")

        alignment, match_rate = self.align(target_ipa, actual_ipa, confidences)
        errors = [p for p in alignment if p.status != "match"]

        return AnalysisResult(
            text=text,
            transcription_confidence=transcription_confidence,
            target_ipa=target_ipa,
            actual_ipa=actual_ipa,
            model_confidence=model_confidence,
            alignment=alignment,
            match_rate=match_rate,
            errors=errors,
        )