data_processing/multimodal_fusion.py

"""Multi-modal data fusion for Myanmar Ghost project.

Fuses audio (prosody) and text to understand sentiment/intensity
in expressions like "ကျေးဇူးပါ" (thank you) which can mean:
- Genuine gratitude (low pitch, slow)
- Sarcasm (high pitch, fast)
- Complaint (negative prosody)
"""

from dataclasses import dataclass
from enum import Enum
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
from torch import Tensor


class SentimentClass(str, Enum):
    """Sentiment classes for thanking expressions."""
    GENUINE = "genuine"          # ရိုးသားခြင်း
    SARCASTIC = "sarcastic"      # သရော်ခြင်း
    COMPLAINING = "complaining"  # မကျေနပ်ခြင်း
    NEUTRAL = "neutral"


@dataclass
class ProsodyFeatures:
    """Prosodic features extracted from audio."""
    mean_pitch: float
    pitch_std: float
    pitch_range: Tuple[float, float]
    mean_energy: float
    energy_std: float
    speaking_rate: float  # syllables per second
    pause_duration: float  # total pause time in seconds
    
    def to_tensor(self) -> Tensor:
        """Convert to PyTorch tensor."""
        return torch.tensor([
            self.mean_pitch,
            self.pitch_std,
            self.pitch_range[0],
            self.pitch_range[1],
            self.mean_energy,
            self.energy_std,
            self.speaking_rate,
            self.pause_duration,
        ], dtype=torch.float32)
    
    def to_dict(self) -> Dict[str, float]:
        """Convert to dictionary."""
        return {
            "mean_pitch": self.mean_pitch,
            "pitch_std": self.pitch_std,
            "pitch_min": self.pitch_range[0],
            "pitch_max": self.pitch_range[1],
            "mean_energy": self.mean_energy,
            "energy_std": self.energy_std,
            "speaking_rate": self.speaking_rate,
            "pause_duration": self.pause_duration,
        }


@dataclass
class TextFeatures:
    """Text-based features for sentiment analysis."""
    text_length: int
    word_count: int
    contains_intensifier: bool  # e.g., "အရမ်း", "များစွာ"
    politeness_level: int  # 1-5 scale
    formality: float  # 0-1 scale
    
    def to_tensor(self) -> Tensor:
        """Convert to PyTorch tensor."""
        return torch.tensor([
            float(self.text_length),
            float(self.word_count),
            float(self.contains_intensifier),
            float(self.politeness_level),
            self.formality,
        ], dtype=torch.float32)


@dataclass
class FusedFeatures:
    """Combined multi-modal features."""
    prosody: ProsodyFeatures
    text: TextFeatures
    sentiment_hint: Optional[SentimentClass] = None
    
    def concat_tensors(self) -> Tensor:
        """Concatenate all features into single tensor."""
        return torch.cat([
            self.prosody.to_tensor(),
            self.text.to_tensor(),
        ])


class ProsodyExtractor:
    """Extract prosodic features from audio."""
    
    # Prosody patterns for different sentiments
    GENUINE_PATTERN = {
        "pitch_range": (50, 200),  # Hz
        "speaking_rate": (2, 4),   # syllables/sec
        "energy_std": (0.1, 0.3),
    }
    
    SARCASTIC_PATTERN = {
        "pitch_range": (200, 400),
        "speaking_rate": (4, 8),
        "energy_std": (0.3, 0.6),
    }
    
    COMPLAINING_PATTERN = {
        "pitch_range": (100, 250),
        "speaking_rate": (3, 6),
        "energy_std": (0.2, 0.5),
    }
    
    def extract_from_audio(
        self,
        audio: np.ndarray,
        sample_rate: int = 16000,
    ) -> ProsodyFeatures:
        """Extract prosodic features from audio signal."""
        import librosa
        
        # Pitch tracking
        pitches, magnitudes = librosa.piptrack(
            y=audio,
            sr=sample_rate,
            n_fft=512,
            hop_length=160,
        )
        
        pitch_values = []
        for i in range(pitches.shape[1]):
            index = magnitudes[:, i].argmax()
            pitch = pitches[index, i]
            if pitch > 0:
                pitch_values.append(pitch)
        
        # Energy
        rms = librosa.feature.rms(y=audio, hop_length=160)[0]
        
        # Speaking rate (syllable detection)
        onsets = librosa.onset.onset_detect(
            y=audio,
            sr=sample_rate,
            hop_length=160,
        )
        
        duration = len(audio) / sample_rate
        speaking_rate = len(onsets) / duration if duration > 0 else 0
        
        # Pause detection
        energy_threshold = np.percentile(rms, 25)
        pauses = rms < energy_threshold
        pause_duration = np.sum(pauses) * 160 / sample_rate
        
        return ProsodyFeatures(
            mean_pitch=np.mean(pitch_values) if pitch_values else 0,
            pitch_std=np.std(pitch_values) if pitch_values else 0,
            pitch_range=(
                np.min(pitch_values) if pitch_values else 0,
                np.max(pitch_values) if pitch_values else 0,
            ),
            mean_energy=np.mean(rms),
            energy_std=np.std(rms),
            speaking_rate=speaking_rate,
            pause_duration=pause_duration,
        )
    
    def infer_sentiment(self, prosody: ProsodyFeatures) -> SentimentClass:
        """Infer sentiment from prosodic features."""
        patterns = [
            (SentimentClass.GENUINE, self.GENUINE_PATTERN),
            (SentimentClass.SARCASTIC, self.SARCASTIC_PATTERN),
            (SentimentClass.COMPLAINING, self.COMPLAINING_PATTERN),
        ]
        
        scores = {}
        for sentiment, pattern in patterns:
            score = 0
            features = prosody.to_dict()
            
            for key, (low, high) in pattern.items():
                if key in features:
                    value = features[key]
                    if low <= value <= high:
                        score += 1
            
            scores[sentiment] = score
        
        return max(scores, key=scores.get)


class TextFeatureExtractor:
    """Extract text-based features."""
    
    INTENSIFIERS = {"အရမ်း", "များစွာ", "ပါး", "သိပ်", "အလွန်"}
    POLITE_WORDS = {"ကျေးဇူး", "�心病", "ဂုဏ်", "အား", "ကြိုးစား", "ပင်ပန်း"}
    
    def extract_from_text(self, text: str) -> TextFeatures:
        """Extract features from text."""
        words = text.split()
        
        has_intensifier = any(
            word in self.INTENSIFIERS for word in words
        )
        
        politeness = self._estimate_politeness(text)
        formality = self._estimate_formality(text)
        
        return TextFeatures(
            text_length=len(text),
            word_count=len(words),
            contains_intensifier=has_intensifier,
            politeness_level=politeness,
            formality=formality,
        )
    
    def _estimate_politeness(self, text: str) -> int:
        """Estimate politeness level (1-5)."""
        score = 3  # default neutral
        polite_count = sum(1 for w in self.POLITE_WORDS if w in text)
        if "ပါ" in text or "ပါး" in text:
            score += 1
        if "ကျေးဇူး" in text:
            score += 1
        if polite_count > 2:
            score += 1
        return min(5, max(1, score))
    
    def _estimate_formality(self, text: str) -> float:
        """Estimate formality (0-1)."""
        formal_markers = {"မှ", "သည်", "ကို", "ဖြင့်", "အား"}
        informal_markers = {"နော်", "ဟုတ်", "မဟုတ်", "လား"}
        
        formal_count = sum(1 for m in formal_markers if m in text)
        informal_count = sum(1 for m in informal_markers if m in text)
        
        if formal_count + informal_count == 0:
            return 0.5
        
        return formal_count / (formal_count + informal_count + 1)


class MultiModalFusion(nn.Module):
    """Fuse audio and text modalities."""
    
    def __init__(
        self,
        prosody_dim: int = 8,
        text_dim: int = 5,
        hidden_dim: int = 64,
        num_classes: int = 4,
    ):
        super().__init__()
        
        self.prosody_encoder = nn.Sequential(
            nn.Linear(prosody_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
        )
        
        self.text_encoder = nn.Sequential(
            nn.Linear(text_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
        )
        
        self.fusion = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden_dim, num_classes),
        )
    
    def forward(self, prosody: Tensor, text: Tensor) -> Tensor:
        """Forward pass."""
        p_encoded = self.prosody_encoder(prosody)
        t_encoded = self.text_encoder(text)
        
        fused = torch.cat([p_encoded, t_encoded], dim=-1)
        logits = self.fusion(fused)
        
        return logits
    
    def predict(self, prosody: Tensor, text: Tensor) -> Tuple[Tensor, Tensor]:
        """Predict sentiment with probabilities."""
        logits = self.forward(prosody, text)
        probs = torch.softmax(logits, dim=-1)
        return logits, probs


class SentimentClassifier:
    """High-level classifier for multi-modal sentiment."""
    
    def __init__(self, model: MultiModalFusion):
        self.model = model
        self.prosody_extractor = ProsodyExtractor()
        self.text_extractor = TextFeatureExtractor()
    
    def classify(
        self,
        audio: np.ndarray,
        text: str,
        return_probs: bool = True,
    ) -> Dict[str, Any]:
        """Classify sentiment from audio and text."""
        prosody_features = self.prosody_extractor.extract_from_audio(audio)
        prosody_hint = self.prosody_extractor.infer_sentiment(prosody_features)
        
        text_features = self.text_extractor.extract_from_text(text)
        
        fused = FusedFeatures(
            prosody=prosody_features,
            text=text_features,
            sentiment_hint=prosody_hint,
        )
        
        prosody_tensor = fused.prosody.to_tensor().unsqueeze(0)
        text_tensor = fused.text.to_tensor().unsqueeze(0)
        
        with torch.no_grad():
            logits, probs = self.model.predict(prosody_tensor, text_tensor)
        
        result = {
            "predicted_class": SentimentClass(probs.argmax().item()).value,
            "prosody_hint": prosody_hint.value,
            "text_features": text_features.to_dict(),
            "prosody_features": prosody_features.to_dict(),
        }
        
        if return_probs:
            result["probabilities"] = {
                c.value: probs[0, i].item()
                for i, c in enumerate(SentimentClass)
            }
        
        return result


def create_fusion_model(
    prosody_dim: int = 8,
    text_dim: int = 5,
    hidden_dim: int = 64,
    num_classes: int = 4,
) -> MultiModalFusion:
    """Factory function to create fusion model."""
    return MultiModalFusion(
        prosody_dim=prosody_dim,
        text_dim=text_dim,
        hidden_dim=hidden_dim,
        num_classes=num_classes,
    )


if __name__ == "__main__":
    # Example usage
    model = create_fusion_model()
    prosody = torch.randn(1, 8)
    text = torch.randn(1, 5)
    
    logits, probs = model.predict(prosody, text)
    print(f"Predicted class: {SentimentClass(probs.argmax().item()).value}")
    print(f"Probabilities: {probs}")