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

//! Prosody prediction for natural speech synthesis
//!
//! This module provides prosody analysis and prediction including:
//! - Pause detection at sentence and phrase boundaries
//! - Emphasis and stress prediction
//! - Intonation pattern prediction (rising/falling)
//! - Speaking rate estimation

use crate::error::Result;
use serde::{Deserialize, Serialize};

/// Configuration for prosody prediction
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProsodyConfig {
    /// Use punctuation for pause prediction
    pub use_punctuation_pauses: bool,
    /// Default speaking rate (1.0 = normal, <1.0 = slower, >1.0 = faster)
    pub default_speaking_rate: f32,
    /// Insert pauses at commas
    pub pause_at_commas: bool,
    /// Insert pauses at sentence boundaries
    pub pause_at_sentences: bool,
    /// Detect emphasis from capitalization
    pub detect_emphasis: bool,
}

impl Default for ProsodyConfig {
    fn default() -> Self {
        Self {
            use_punctuation_pauses: true,
            default_speaking_rate: 1.0,
            pause_at_commas: true,
            pause_at_sentences: true,
            detect_emphasis: true,
        }
    }
}

/// Prosody marker types
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum ProsodyMarker {
    /// Pause with duration in milliseconds
    Pause(u32),
    /// Emphasis on word or syllable
    Emphasis,
    /// Rising intonation (questions)
    Rising,
    /// Falling intonation (statements)
    Falling,
    /// Continuation (phrase boundary)
    Continuation,
}

/// Prosody information for a text segment
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProsodyAnnotation {
    /// Text segment
    pub text: String,
    /// Start position in original text
    pub start_pos: usize,
    /// End position in original text
    pub end_pos: usize,
    /// Prosody markers for this segment
    pub markers: Vec<ProsodyMarker>,
    /// Duration factor (relative to normal speaking rate)
    pub duration_factor: f32,
}

impl ProsodyAnnotation {
    /// Create a new prosody annotation
    #[must_use]
    pub fn new(text: String, start_pos: usize, end_pos: usize) -> Self {
        Self {
            text,
            start_pos,
            end_pos,
            markers: Vec::new(),
            duration_factor: 1.0,
        }
    }

    /// Add a marker to this annotation
    pub fn add_marker(&mut self, marker: ProsodyMarker) {
        self.markers.push(marker);
    }

    /// Check if annotation has a specific marker type
    #[must_use]
    pub fn has_marker(&self, marker_type: &ProsodyMarker) -> bool {
        self.markers
            .iter()
            .any(|m| std::mem::discriminant(m) == std::mem::discriminant(marker_type))
    }
}

/// Prosody predictor
#[derive(Debug)]
pub struct Prosody {
    config: ProsodyConfig,
}

impl Prosody {
    /// Create a new prosody predictor with the given configuration
    #[must_use]
    pub fn new(config: ProsodyConfig) -> Self {
        Self { config }
    }

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

    /// Predict prosody for the given text
    ///
    /// # Errors
    ///
    /// Returns error if prosody prediction fails
    pub fn predict(&self, text: &str) -> Result<Vec<ProsodyAnnotation>> {
        let mut annotations = Vec::new();

        // Split into sentences
        let sentences = self.detect_sentence_boundaries(text)?;

        for (sentence, start_pos) in sentences {
            let sentence_annotations = self.process_sentence(&sentence, start_pos)?;
            annotations.extend(sentence_annotations);
        }

        Ok(annotations)
    }

    /// Detect sentence boundaries
    fn detect_sentence_boundaries(&self, text: &str) -> Result<Vec<(String, usize)>> {
        let mut sentences = Vec::new();
        let mut current_sentence = String::new();
        let mut start_pos = 0;

        for (i, ch) in text.char_indices() {
            current_sentence.push(ch);

            // Check for sentence-ending punctuation
            if ch == '.' || ch == '!' || ch == '?' {
                // Look ahead to see if this is really the end of a sentence
                let next_char = text.chars().nth(i + 1);
                if next_char.is_none() || next_char == Some(' ') || next_char == Some('\n') {
                    sentences.push((current_sentence.trim().to_string(), start_pos));
                    current_sentence.clear();
                    start_pos = i + 1;
                }
            }
        }

        // Add remaining text as a sentence
        if !current_sentence.trim().is_empty() {
            sentences.push((current_sentence.trim().to_string(), start_pos));
        }

        Ok(sentences)
    }

    /// Process a single sentence
    fn process_sentence(&self, sentence: &str, base_pos: usize) -> Result<Vec<ProsodyAnnotation>> {
        let mut annotations = Vec::new();

        // Detect phrase boundaries (commas, semicolons, etc.)
        let phrases = self.detect_phrase_boundaries(sentence)?;

        for (i, (phrase, start_offset)) in phrases.iter().enumerate() {
            let start_pos = base_pos + start_offset;
            let end_pos = start_pos + phrase.len();

            let mut annotation = ProsodyAnnotation::new(phrase.clone(), start_pos, end_pos);

            // Add pause markers
            if self.config.use_punctuation_pauses {
                // Pause at phrase boundaries
                if i < phrases.len() - 1 && self.config.pause_at_commas {
                    annotation.add_marker(ProsodyMarker::Pause(200));
                }
            }

            // Detect emphasis from capitalization
            if self.config.detect_emphasis {
                if phrase.chars().filter(|c| c.is_uppercase()).count() > phrase.len() / 2 {
                    annotation.add_marker(ProsodyMarker::Emphasis);
                }
            }

            // Detect intonation from punctuation
            if sentence.ends_with('?') && i == phrases.len() - 1 {
                annotation.add_marker(ProsodyMarker::Rising);
            } else if sentence.ends_with('.') || sentence.ends_with('!') {
                if i == phrases.len() - 1 {
                    annotation.add_marker(ProsodyMarker::Falling);
                } else {
                    annotation.add_marker(ProsodyMarker::Continuation);
                }
            }

            // Estimate duration factor
            annotation.duration_factor = self.estimate_duration_factor(phrase);

            annotations.push(annotation);
        }

        // Add sentence-final pause
        if self.config.pause_at_sentences && !annotations.is_empty() {
            if let Some(last) = annotations.last_mut() {
                last.add_marker(ProsodyMarker::Pause(400));
            }
        }

        Ok(annotations)
    }

    /// Detect phrase boundaries within a sentence
    fn detect_phrase_boundaries(&self, sentence: &str) -> Result<Vec<(String, usize)>> {
        let mut phrases = Vec::new();
        let mut current_phrase = String::new();
        let mut start_pos = 0;

        for (i, ch) in sentence.char_indices() {
            current_phrase.push(ch);

            // Check for phrase-ending punctuation
            if ch == ',' || ch == ';' || ch == ':' {
                phrases.push((current_phrase.trim().to_string(), start_pos));
                current_phrase.clear();
                start_pos = i + 1;
            }
        }

        // Add remaining text as a phrase
        if !current_phrase.trim().is_empty() {
            phrases.push((current_phrase.trim().to_string(), start_pos));
        }

        Ok(phrases)
    }

    /// Estimate duration factor for a phrase
    fn estimate_duration_factor(&self, phrase: &str) -> f32 {
        // Base duration on speaking rate
        let mut factor = self.config.default_speaking_rate;

        // Longer phrases might be spoken slightly slower
        let word_count = phrase.split_whitespace().count();
        if word_count > 10 {
            factor *= 0.95;
        }

        // Short exclamations might be faster
        if word_count <= 2 && phrase.contains('!') {
            factor *= 1.1;
        }

        factor
    }

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

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

    #[test]
    fn test_default_config() {
        let config = ProsodyConfig::default();
        assert!(config.use_punctuation_pauses);
        assert_eq!(config.default_speaking_rate, 1.0);
        assert!(config.pause_at_commas);
        assert!(config.pause_at_sentences);
        assert!(config.detect_emphasis);
    }

    #[test]
    fn test_prosody_marker_equality() {
        let pause1 = ProsodyMarker::Pause(200);
        let pause2 = ProsodyMarker::Pause(300);
        let emphasis = ProsodyMarker::Emphasis;

        assert_eq!(pause1, ProsodyMarker::Pause(200));
        assert_ne!(pause1, pause2);
        assert_ne!(pause1, emphasis);
    }

    #[test]
    fn test_prosody_annotation_creation() {
        let mut annotation = ProsodyAnnotation::new("hello".to_string(), 0, 5);
        assert_eq!(annotation.text, "hello");
        assert_eq!(annotation.start_pos, 0);
        assert_eq!(annotation.end_pos, 5);
        assert!(annotation.markers.is_empty());
        assert_eq!(annotation.duration_factor, 1.0);

        annotation.add_marker(ProsodyMarker::Emphasis);
        assert_eq!(annotation.markers.len(), 1);
        assert!(annotation.has_marker(&ProsodyMarker::Emphasis));
    }

    #[test]
    fn test_prosody_creation() {
        let prosody = Prosody::default();
        assert_eq!(prosody.config().default_speaking_rate, 1.0);
    }

    #[test]
    fn test_sentence_boundary_detection() {
        let prosody = Prosody::default();
        let text = "Hello world. How are you?";
        let sentences = prosody.detect_sentence_boundaries(text).unwrap();

        assert_eq!(sentences.len(), 2);
        assert_eq!(sentences[0].0, "Hello world.");
        assert_eq!(sentences[1].0, "How are you?");
    }

    #[test]
    fn test_phrase_boundary_detection() {
        let prosody = Prosody::default();
        let sentence = "Hello, how are you today?";
        let phrases = prosody.detect_phrase_boundaries(sentence).unwrap();

        assert_eq!(phrases.len(), 2);
        assert_eq!(phrases[0].0, "Hello,");
        assert_eq!(phrases[1].0, "how are you today?");
    }

    #[test]
    fn test_simple_sentence_prosody() {
        let prosody = Prosody::default();
        let result = prosody.predict("Hello world.").unwrap();

        assert!(!result.is_empty());
        // Should have a sentence-final pause
        assert!(result.last().unwrap().has_marker(&ProsodyMarker::Pause(0)));
    }

    #[test]
    fn test_question_intonation() {
        let prosody = Prosody::default();
        let result = prosody.predict("How are you?").unwrap();

        assert!(!result.is_empty());
        // Should have rising intonation
        assert!(result.last().unwrap().has_marker(&ProsodyMarker::Rising));
    }

    #[test]
    fn test_statement_intonation() {
        let prosody = Prosody::default();
        let result = prosody.predict("I am fine.").unwrap();

        assert!(!result.is_empty());
        // Should have falling intonation
        assert!(result.last().unwrap().has_marker(&ProsodyMarker::Falling));
    }

    #[test]
    fn test_comma_pause() {
        let config = ProsodyConfig {
            use_punctuation_pauses: true,
            pause_at_commas: true,
            ..Default::default()
        };
        let prosody = Prosody::new(config);
        let result = prosody.predict("Hello, world.").unwrap();

        assert_eq!(result.len(), 2);
        // First phrase should have a pause after comma
        assert!(result[0].has_marker(&ProsodyMarker::Pause(0)));
    }

    #[test]
    fn test_multiple_sentences() {
        let prosody = Prosody::default();
        let result = prosody
            .predict("Hello world. How are you? I am fine.")
            .unwrap();

        assert!(!result.is_empty());
        // Multiple sentences should be detected
        assert!(result.len() >= 3);
    }

    #[test]
    fn test_emphasis_detection() {
        let config = ProsodyConfig {
            detect_emphasis: true,
            ..Default::default()
        };
        let prosody = Prosody::new(config);
        let result = prosody.predict("HELLO WORLD").unwrap();

        assert!(!result.is_empty());
        // Should detect emphasis from capitalization
        assert!(result[0].has_marker(&ProsodyMarker::Emphasis));
    }

    #[test]
    fn test_duration_factor_estimation() {
        let prosody = Prosody::default();

        let short_factor = prosody.estimate_duration_factor("Hi!");
        let normal_factor = prosody.estimate_duration_factor("This is a normal sentence");
        let long_factor = prosody.estimate_duration_factor(
            "This is a very long sentence with many words that should be spoken slowly",
        );

        assert!(short_factor >= normal_factor);
        assert!(normal_factor >= long_factor);
    }

    #[test]
    fn test_no_punctuation_pauses() {
        let config = ProsodyConfig {
            use_punctuation_pauses: false,
            pause_at_commas: false,
            pause_at_sentences: false,
            ..Default::default()
        };
        let prosody = Prosody::new(config);
        let result = prosody.predict("Hello, world.").unwrap();

        // Should not have pauses
        for annotation in &result {
            assert!(!annotation.has_marker(&ProsodyMarker::Pause(0)));
        }
    }

    #[test]
    fn test_speaking_rate() {
        let config = ProsodyConfig {
            default_speaking_rate: 1.5,
            ..Default::default()
        };
        let prosody = Prosody::new(config);
        let result = prosody.predict("Hello world.").unwrap();

        assert!(!result.is_empty());
        // Duration factors should reflect faster speaking rate
        assert!(result[0].duration_factor >= 1.0);
    }

    #[test]
    fn test_semicolon_boundary() {
        let prosody = Prosody::default();
        let result = prosody.predict("First part; second part.").unwrap();

        assert_eq!(result.len(), 2);
        assert_eq!(result[0].text, "First part;");
        assert_eq!(result[1].text, "second part.");
    }

    #[test]
    fn test_exclamation_intonation() {
        let prosody = Prosody::default();
        let result = prosody.predict("Hello!").unwrap();

        assert!(!result.is_empty());
        // Exclamations should have falling intonation
        assert!(result.last().unwrap().has_marker(&ProsodyMarker::Falling));
    }

    #[test]
    fn test_empty_text() {
        let prosody = Prosody::default();
        let result = prosody.predict("").unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_whitespace_only() {
        let prosody = Prosody::default();
        let result = prosody.predict("   ").unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_complex_punctuation() {
        let prosody = Prosody::default();
        let result = prosody.predict("Wait... what? Really!").unwrap();

        assert!(!result.is_empty());
        // Should handle multiple punctuation marks
        assert!(result.len() >= 2);
    }
}