rustytorch/crates/models/rtx-tts/src/text/phonemizer.rs

//! Grapheme-to-phoneme conversion for TTS
//!
//! This module provides phonemization capabilities using:
//! - Dictionary-based lookup for common words
//! - Rule-based fallback for unknown words
//! - ARPAbet phoneme representation for English

use crate::error::{Result, TtsError};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Phonemizer backend type
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum PhonemizerBackend {
    /// Dictionary-based lookup
    Dictionary,
    /// Rule-based phonemization
    Rule,
    /// Grapheme-to-phoneme model (future)
    G2P,
}

/// Configuration for phonemization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PhonemizerConfig {
    /// Language code (e.g., "en-us", "en-gb")
    pub language: String,
    /// Phonemization backend
    pub backend: PhonemizerBackend,
    /// Include stress markers (0=no stress, 1=primary, 2=secondary)
    pub include_stress: bool,
}

impl Default for PhonemizerConfig {
    fn default() -> Self {
        Self {
            language: "en-us".to_string(),
            backend: PhonemizerBackend::Dictionary,
            include_stress: true,
        }
    }
}

/// Phoneme representation with optional stress
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Phoneme {
    /// Phoneme symbol (ARPAbet for English)
    pub symbol: String,
    /// Stress level (0=unstressed, 1=primary, 2=secondary)
    pub stress: Option<u8>,
}

impl Phoneme {
    /// Create a new phoneme
    #[must_use]
    pub fn new(symbol: String, stress: Option<u8>) -> Self {
        Self { symbol, stress }
    }

    /// Create a phoneme without stress
    #[must_use]
    pub fn unstressed(symbol: String) -> Self {
        Self {
            symbol,
            stress: None,
        }
    }

    /// Convert phoneme to string representation
    #[must_use]
    pub fn to_string_with_stress(&self) -> String {
        if let Some(stress) = self.stress {
            format!("{}{}", self.symbol, stress)
        } else {
            self.symbol.clone()
        }
    }
}

/// Grapheme-to-phoneme converter
#[derive(Debug)]
pub struct Phonemizer {
    config: PhonemizerConfig,
    dictionary: HashMap<String, Vec<Phoneme>>,
}

impl Phonemizer {
    /// Create a new phonemizer with the given configuration
    #[must_use]
    pub fn new(config: PhonemizerConfig) -> Self {
        let mut dictionary = HashMap::new();

        // Build English pronunciation dictionary (ARPAbet)
        // Common words with stress markers
        Self::populate_dictionary(&mut dictionary);

        Self { config, dictionary }
    }

    /// Create a phonemizer with default configuration
    #[must_use]
    pub fn default() -> Self {
        Self::new(PhonemizerConfig::default())
    }

    /// Get the configuration
    #[must_use]
    pub fn config(&self) -> &PhonemizerConfig {
        &self.config
    }

    /// Populate the pronunciation dictionary with common English words
    fn populate_dictionary(dict: &mut HashMap<String, Vec<Phoneme>>) {
        // Pronouns
        dict.insert(
            "i".to_string(),
            vec![Phoneme::new("AY".to_string(), Some(1))],
        );
        dict.insert(
            "you".to_string(),
            vec![
                Phoneme::new("Y".to_string(), None),
                Phoneme::new("UW".to_string(), Some(1)),
            ],
        );
        dict.insert(
            "he".to_string(),
            vec![
                Phoneme::new("HH".to_string(), None),
                Phoneme::new("IY".to_string(), Some(1)),
            ],
        );
        dict.insert(
            "she".to_string(),
            vec![
                Phoneme::new("SH".to_string(), None),
                Phoneme::new("IY".to_string(), Some(1)),
            ],
        );
        dict.insert(
            "it".to_string(),
            vec![
                Phoneme::new("IH".to_string(), Some(1)),
                Phoneme::new("T".to_string(), None),
            ],
        );
        dict.insert(
            "we".to_string(),
            vec![
                Phoneme::new("W".to_string(), None),
                Phoneme::new("IY".to_string(), Some(1)),
            ],
        );
        dict.insert(
            "they".to_string(),
            vec![
                Phoneme::new("DH".to_string(), None),
                Phoneme::new("EY".to_string(), Some(1)),
            ],
        );

        // Common verbs
        dict.insert(
            "is".to_string(),
            vec![
                Phoneme::new("IH".to_string(), Some(1)),
                Phoneme::new("Z".to_string(), None),
            ],
        );
        dict.insert(
            "are".to_string(),
            vec![
                Phoneme::new("AA".to_string(), Some(1)),
                Phoneme::new("R".to_string(), None),
            ],
        );
        dict.insert(
            "was".to_string(),
            vec![
                Phoneme::new("W".to_string(), None),
                Phoneme::new("AA".to_string(), Some(1)),
                Phoneme::new("Z".to_string(), None),
            ],
        );
        dict.insert(
            "have".to_string(),
            vec![
                Phoneme::new("HH".to_string(), None),
                Phoneme::new("AE".to_string(), Some(1)),
                Phoneme::new("V".to_string(), None),
            ],
        );
        dict.insert(
            "had".to_string(),
            vec![
                Phoneme::new("HH".to_string(), None),
                Phoneme::new("AE".to_string(), Some(1)),
                Phoneme::new("D".to_string(), None),
            ],
        );
        dict.insert(
            "do".to_string(),
            vec![
                Phoneme::new("D".to_string(), None),
                Phoneme::new("UW".to_string(), Some(1)),
            ],
        );
        dict.insert(
            "did".to_string(),
            vec![
                Phoneme::new("D".to_string(), None),
                Phoneme::new("IH".to_string(), Some(1)),
                Phoneme::new("D".to_string(), None),
            ],
        );
        dict.insert(
            "said".to_string(),
            vec![
                Phoneme::new("S".to_string(), None),
                Phoneme::new("EH".to_string(), Some(1)),
                Phoneme::new("D".to_string(), None),
            ],
        );
        dict.insert(
            "go".to_string(),
            vec![
                Phoneme::new("G".to_string(), None),
                Phoneme::new("OW".to_string(), Some(1)),
            ],
        );
        dict.insert(
            "went".to_string(),
            vec![
                Phoneme::new("W".to_string(), None),
                Phoneme::new("EH".to_string(), Some(1)),
                Phoneme::new("N".to_string(), None),
                Phoneme::new("T".to_string(), None),
            ],
        );
        dict.insert(
            "make".to_string(),
            vec![
                Phoneme::new("M".to_string(), None),
                Phoneme::new("EY".to_string(), Some(1)),
                Phoneme::new("K".to_string(), None),
            ],
        );
        dict.insert(
            "get".to_string(),
            vec![
                Phoneme::new("G".to_string(), None),
                Phoneme::new("EH".to_string(), Some(1)),
                Phoneme::new("T".to_string(), None),
            ],
        );

        // Common nouns
        dict.insert(
            "time".to_string(),
            vec![
                Phoneme::new("T".to_string(), None),
                Phoneme::new("AY".to_string(), Some(1)),
                Phoneme::new("M".to_string(), None),
            ],
        );
        dict.insert(
            "day".to_string(),
            vec![
                Phoneme::new("D".to_string(), None),
                Phoneme::new("EY".to_string(), Some(1)),
            ],
        );
        dict.insert(
            "man".to_string(),
            vec![
                Phoneme::new("M".to_string(), None),
                Phoneme::new("AE".to_string(), Some(1)),
                Phoneme::new("N".to_string(), None),
            ],
        );
        dict.insert(
            "woman".to_string(),
            vec![
                Phoneme::new("W".to_string(), None),
                Phoneme::new("UH".to_string(), Some(1)),
                Phoneme::new("M".to_string(), None),
                Phoneme::new("AH".to_string(), Some(0)),
                Phoneme::new("N".to_string(), None),
            ],
        );
        dict.insert(
            "child".to_string(),
            vec![
                Phoneme::new("CH".to_string(), None),
                Phoneme::new("AY".to_string(), Some(1)),
                Phoneme::new("L".to_string(), None),
                Phoneme::new("D".to_string(), None),
            ],
        );
        dict.insert(
            "world".to_string(),
            vec![
                Phoneme::new("W".to_string(), None),
                Phoneme::new("ER".to_string(), Some(1)),
                Phoneme::new("L".to_string(), None),
                Phoneme::new("D".to_string(), None),
            ],
        );
        dict.insert(
            "house".to_string(),
            vec![
                Phoneme::new("HH".to_string(), None),
                Phoneme::new("AW".to_string(), Some(1)),
                Phoneme::new("S".to_string(), None),
            ],
        );
        dict.insert(
            "book".to_string(),
            vec![
                Phoneme::new("B".to_string(), None),
                Phoneme::new("UH".to_string(), Some(1)),
                Phoneme::new("K".to_string(), None),
            ],
        );

        // Common adjectives
        dict.insert(
            "good".to_string(),
            vec![
                Phoneme::new("G".to_string(), None),
                Phoneme::new("UH".to_string(), Some(1)),
                Phoneme::new("D".to_string(), None),
            ],
        );
        dict.insert(
            "great".to_string(),
            vec![
                Phoneme::new("G".to_string(), None),
                Phoneme::new("R".to_string(), None),
                Phoneme::new("EY".to_string(), Some(1)),
                Phoneme::new("T".to_string(), None),
            ],
        );
        dict.insert(
            "new".to_string(),
            vec![
                Phoneme::new("N".to_string(), None),
                Phoneme::new("UW".to_string(), Some(1)),
            ],
        );
        dict.insert(
            "old".to_string(),
            vec![
                Phoneme::new("OW".to_string(), Some(1)),
                Phoneme::new("L".to_string(), None),
                Phoneme::new("D".to_string(), None),
            ],
        );
        dict.insert(
            "big".to_string(),
            vec![
                Phoneme::new("B".to_string(), None),
                Phoneme::new("IH".to_string(), Some(1)),
                Phoneme::new("G".to_string(), None),
            ],
        );
        dict.insert(
            "small".to_string(),
            vec![
                Phoneme::new("S".to_string(), None),
                Phoneme::new("M".to_string(), None),
                Phoneme::new("AO".to_string(), Some(1)),
                Phoneme::new("L".to_string(), None),
            ],
        );

        // Articles and common words
        dict.insert(
            "the".to_string(),
            vec![
                Phoneme::new("DH".to_string(), None),
                Phoneme::new("AH".to_string(), Some(0)),
            ],
        );
        dict.insert(
            "a".to_string(),
            vec![Phoneme::new("AH".to_string(), Some(0))],
        );
        dict.insert(
            "an".to_string(),
            vec![
                Phoneme::new("AE".to_string(), Some(1)),
                Phoneme::new("N".to_string(), None),
            ],
        );
        dict.insert(
            "and".to_string(),
            vec![
                Phoneme::new("AE".to_string(), Some(1)),
                Phoneme::new("N".to_string(), None),
                Phoneme::new("D".to_string(), None),
            ],
        );
        dict.insert(
            "or".to_string(),
            vec![
                Phoneme::new("AO".to_string(), Some(1)),
                Phoneme::new("R".to_string(), None),
            ],
        );
        dict.insert(
            "but".to_string(),
            vec![
                Phoneme::new("B".to_string(), None),
                Phoneme::new("AH".to_string(), Some(1)),
                Phoneme::new("T".to_string(), None),
            ],
        );
        dict.insert(
            "to".to_string(),
            vec![
                Phoneme::new("T".to_string(), None),
                Phoneme::new("UW".to_string(), Some(1)),
            ],
        );
        dict.insert(
            "in".to_string(),
            vec![
                Phoneme::new("IH".to_string(), Some(1)),
                Phoneme::new("N".to_string(), None),
            ],
        );
        dict.insert(
            "on".to_string(),
            vec![
                Phoneme::new("AA".to_string(), Some(1)),
                Phoneme::new("N".to_string(), None),
            ],
        );
        dict.insert(
            "at".to_string(),
            vec![
                Phoneme::new("AE".to_string(), Some(1)),
                Phoneme::new("T".to_string(), None),
            ],
        );
        dict.insert(
            "for".to_string(),
            vec![
                Phoneme::new("F".to_string(), None),
                Phoneme::new("AO".to_string(), Some(1)),
                Phoneme::new("R".to_string(), None),
            ],
        );
        dict.insert(
            "with".to_string(),
            vec![
                Phoneme::new("W".to_string(), None),
                Phoneme::new("IH".to_string(), Some(1)),
                Phoneme::new("DH".to_string(), None),
            ],
        );

        // Numbers
        dict.insert(
            "one".to_string(),
            vec![
                Phoneme::new("W".to_string(), None),
                Phoneme::new("AH".to_string(), Some(1)),
                Phoneme::new("N".to_string(), None),
            ],
        );
        dict.insert(
            "two".to_string(),
            vec![
                Phoneme::new("T".to_string(), None),
                Phoneme::new("UW".to_string(), Some(1)),
            ],
        );
        dict.insert(
            "three".to_string(),
            vec![
                Phoneme::new("TH".to_string(), None),
                Phoneme::new("R".to_string(), None),
                Phoneme::new("IY".to_string(), Some(1)),
            ],
        );
        dict.insert(
            "four".to_string(),
            vec![
                Phoneme::new("F".to_string(), None),
                Phoneme::new("AO".to_string(), Some(1)),
                Phoneme::new("R".to_string(), None),
            ],
        );
        dict.insert(
            "five".to_string(),
            vec![
                Phoneme::new("F".to_string(), None),
                Phoneme::new("AY".to_string(), Some(1)),
                Phoneme::new("V".to_string(), None),
            ],
        );

        // Common contractions
        dict.insert(
            "can't".to_string(),
            vec![
                Phoneme::new("K".to_string(), None),
                Phoneme::new("AE".to_string(), Some(1)),
                Phoneme::new("N".to_string(), None),
                Phoneme::new("T".to_string(), None),
            ],
        );
        dict.insert(
            "don't".to_string(),
            vec![
                Phoneme::new("D".to_string(), None),
                Phoneme::new("OW".to_string(), Some(1)),
                Phoneme::new("N".to_string(), None),
                Phoneme::new("T".to_string(), None),
            ],
        );
        dict.insert(
            "won't".to_string(),
            vec![
                Phoneme::new("W".to_string(), None),
                Phoneme::new("OW".to_string(), Some(1)),
                Phoneme::new("N".to_string(), None),
                Phoneme::new("T".to_string(), None),
            ],
        );

        // Greetings
        dict.insert(
            "hello".to_string(),
            vec![
                Phoneme::new("HH".to_string(), None),
                Phoneme::new("AH".to_string(), Some(0)),
                Phoneme::new("L".to_string(), None),
                Phoneme::new("OW".to_string(), Some(1)),
            ],
        );
        dict.insert(
            "hi".to_string(),
            vec![
                Phoneme::new("HH".to_string(), None),
                Phoneme::new("AY".to_string(), Some(1)),
            ],
        );
        dict.insert(
            "goodbye".to_string(),
            vec![
                Phoneme::new("G".to_string(), None),
                Phoneme::new("UH".to_string(), Some(1)),
                Phoneme::new("D".to_string(), None),
                Phoneme::new("B".to_string(), None),
                Phoneme::new("AY".to_string(), Some(1)),
            ],
        );
    }

    /// Convert text to phonemes
    ///
    /// # Errors
    ///
    /// Returns error if phonemization fails
    pub fn g2p(&self, text: &str) -> Result<Vec<Vec<Phoneme>>> {
        let words: Vec<&str> = text.split_whitespace().collect();
        let mut result = Vec::new();

        for word in words {
            let word_lower = word.to_lowercase();
            let phonemes = self.phonemize_word(&word_lower)?;
            result.push(phonemes);
        }

        Ok(result)
    }

    /// Phonemize a single word
    fn phonemize_word(&self, word: &str) -> Result<Vec<Phoneme>> {
        // Try dictionary lookup first
        if let Some(phonemes) = self.dictionary.get(word) {
            return Ok(phonemes.clone());
        }

        // Fallback to rule-based phonemization
        self.rule_based_phonemization(word)
    }

    /// Rule-based phonemization for unknown words
    fn rule_based_phonemization(&self, word: &str) -> Result<Vec<Phoneme>> {
        let mut phonemes = Vec::new();
        let chars: Vec<char> = word.chars().collect();
        let mut i = 0;

        while i < chars.len() {
            let ch = chars[i];

            // Simple rule-based conversion (very basic)
            let phoneme = match ch {
                'a' => Phoneme::new("AE".to_string(), Some(1)),
                'b' => Phoneme::unstressed("B".to_string()),
                'c' => Phoneme::unstressed("K".to_string()),
                'd' => Phoneme::unstressed("D".to_string()),
                'e' => Phoneme::new("EH".to_string(), Some(1)),
                'f' => Phoneme::unstressed("F".to_string()),
                'g' => Phoneme::unstressed("G".to_string()),
                'h' => Phoneme::unstressed("HH".to_string()),
                'i' => Phoneme::new("IH".to_string(), Some(1)),
                'j' => Phoneme::unstressed("JH".to_string()),
                'k' => Phoneme::unstressed("K".to_string()),
                'l' => Phoneme::unstressed("L".to_string()),
                'm' => Phoneme::unstressed("M".to_string()),
                'n' => Phoneme::unstressed("N".to_string()),
                'o' => Phoneme::new("AA".to_string(), Some(1)),
                'p' => Phoneme::unstressed("P".to_string()),
                'q' => Phoneme::unstressed("K".to_string()),
                'r' => Phoneme::unstressed("R".to_string()),
                's' => Phoneme::unstressed("S".to_string()),
                't' => Phoneme::unstressed("T".to_string()),
                'u' => Phoneme::new("AH".to_string(), Some(1)),
                'v' => Phoneme::unstressed("V".to_string()),
                'w' => Phoneme::unstressed("W".to_string()),
                'x' => Phoneme::unstressed("K".to_string()),
                'y' => Phoneme::unstressed("Y".to_string()),
                'z' => Phoneme::unstressed("Z".to_string()),
                _ => {
                    i += 1;
                    continue;
                }
            };

            phonemes.push(phoneme);
            i += 1;
        }

        if phonemes.is_empty() {
            return Err(TtsError::Phonemization(format!(
                "Could not phonemize word: {word}"
            )));
        }

        Ok(phonemes)
    }

    /// Convert phonemes to string representation
    #[must_use]
    pub fn phonemes_to_string(&self, phonemes: &[Vec<Phoneme>]) -> String {
        phonemes
            .iter()
            .map(|word_phonemes| {
                word_phonemes
                    .iter()
                    .map(|p| {
                        if self.config.include_stress {
                            p.to_string_with_stress()
                        } else {
                            p.symbol.clone()
                        }
                    })
                    .collect::<Vec<_>>()
                    .join(" ")
            })
            .collect::<Vec<_>>()
            .join(" | ")
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_default_config() {
        let config = PhonemizerConfig::default();
        assert_eq!(config.language, "en-us");
        assert_eq!(config.backend, PhonemizerBackend::Dictionary);
        assert!(config.include_stress);
    }

    #[test]
    fn test_phoneme_creation() {
        let phoneme = Phoneme::new("AE".to_string(), Some(1));
        assert_eq!(phoneme.symbol, "AE");
        assert_eq!(phoneme.stress, Some(1));
        assert_eq!(phoneme.to_string_with_stress(), "AE1");

        let unstressed = Phoneme::unstressed("B".to_string());
        assert_eq!(unstressed.symbol, "B");
        assert_eq!(unstressed.stress, None);
        assert_eq!(unstressed.to_string_with_stress(), "B");
    }

    #[test]
    fn test_phonemizer_creation() {
        let phonemizer = Phonemizer::default();
        assert!(!phonemizer.dictionary.is_empty());
    }

    #[test]
    fn test_dictionary_lookup() {
        let phonemizer = Phonemizer::default();
        let result = phonemizer.g2p("hello").unwrap();
        assert_eq!(result.len(), 1);
        assert!(!result[0].is_empty());
    }

    #[test]
    fn test_common_words() {
        let phonemizer = Phonemizer::default();

        let result = phonemizer.g2p("the").unwrap();
        assert_eq!(result.len(), 1);
        assert_eq!(result[0][0].symbol, "DH");

        let result = phonemizer.g2p("hello").unwrap();
        assert_eq!(result.len(), 1);
        assert_eq!(result[0][0].symbol, "HH");
    }

    #[test]
    fn test_multiple_words() {
        let phonemizer = Phonemizer::default();
        let result = phonemizer.g2p("hello world").unwrap();
        assert_eq!(result.len(), 2);
        assert!(!result[0].is_empty());
        assert!(!result[1].is_empty());
    }

    #[test]
    fn test_pronouns() {
        let phonemizer = Phonemizer::default();

        let result = phonemizer.g2p("i").unwrap();
        assert_eq!(result[0][0].symbol, "AY");
        assert_eq!(result[0][0].stress, Some(1));

        let result = phonemizer.g2p("you").unwrap();
        assert_eq!(result[0][0].symbol, "Y");
        assert_eq!(result[0][1].symbol, "UW");
    }

    #[test]
    fn test_verbs() {
        let phonemizer = Phonemizer::default();

        let result = phonemizer.g2p("is").unwrap();
        assert_eq!(result[0][0].symbol, "IH");
        assert_eq!(result[0][1].symbol, "Z");

        let result = phonemizer.g2p("have").unwrap();
        assert_eq!(result[0].len(), 3);
    }

    #[test]
    fn test_numbers() {
        let phonemizer = Phonemizer::default();

        let result = phonemizer.g2p("one two three").unwrap();
        assert_eq!(result.len(), 3);
        assert!(!result[0].is_empty());
        assert!(!result[1].is_empty());
        assert!(!result[2].is_empty());
    }

    #[test]
    fn test_contractions() {
        let phonemizer = Phonemizer::default();

        let result = phonemizer.g2p("can't").unwrap();
        assert!(!result[0].is_empty());

        let result = phonemizer.g2p("don't").unwrap();
        assert!(!result[0].is_empty());
    }

    #[test]
    fn test_rule_based_fallback() {
        let phonemizer = Phonemizer::default();

        // Unknown word should use rule-based phonemization
        let result = phonemizer.g2p("xyz").unwrap();
        assert_eq!(result.len(), 1);
        assert!(!result[0].is_empty());
    }

    #[test]
    fn test_phonemes_to_string_with_stress() {
        let phonemizer = Phonemizer::default();
        let result = phonemizer.g2p("hello").unwrap();
        let phoneme_str = phonemizer.phonemes_to_string(&result);
        assert!(phoneme_str.contains("HH"));
    }

    #[test]
    fn test_phonemes_to_string_without_stress() {
        let config = PhonemizerConfig {
            language: "en-us".to_string(),
            backend: PhonemizerBackend::Dictionary,
            include_stress: false,
        };
        let phonemizer = Phonemizer::new(config);
        let result = phonemizer.g2p("hello").unwrap();
        let phoneme_str = phonemizer.phonemes_to_string(&result);
        assert!(phoneme_str.contains("HH"));
        assert!(!phoneme_str.contains('0'));
        assert!(!phoneme_str.contains('1'));
    }

    #[test]
    fn test_case_insensitive() {
        let phonemizer = Phonemizer::default();

        let result1 = phonemizer.g2p("HELLO").unwrap();
        let result2 = phonemizer.g2p("hello").unwrap();
        let result3 = phonemizer.g2p("Hello").unwrap();

        assert_eq!(result1, result2);
        assert_eq!(result2, result3);
    }

    #[test]
    fn test_stress_markers() {
        let phonemizer = Phonemizer::default();
        let result = phonemizer.g2p("hello").unwrap();

        // Check that stress markers are present
        let has_stress = result[0].iter().any(|p| p.stress.is_some());
        assert!(has_stress);
    }

    #[test]
    fn test_sentence_phonemization() {
        let phonemizer = Phonemizer::default();
        let result = phonemizer.g2p("the quick brown fox").unwrap();
        assert_eq!(result.len(), 4);

        // Verify each word has phonemes
        for word_phonemes in &result {
            assert!(!word_phonemes.is_empty());
        }
    }
}