"""Multi-modal sentiment model combining audio and text."""

import logging
from typing import Any, Dict, Optional

import torch
import torch.nn as nn

from .base_model import BaseModel

logger = logging.getLogger(__name__)


class MultiModalSentimentModel(BaseModel):
    """Multi-modal model combining text and audio features."""
    
    def __init__(
        self,
        text_dim: int = 768,
        audio_dim: int = 8,
        hidden_dim: int = 256,
        num_classes: int = 4,
        dropout: float = 0.2,
    ):
        """
        Args:
            text_dim: Text embedding dimension
            audio_dim: Audio feature dimension
            hidden_dim: Hidden layer dimension
            num_classes: Number of sentiment classes
            dropout: Dropout rate
        """
        super().__init__()
        
        self.text_dim = text_dim
        self.audio_dim = audio_dim
        
        # Text encoder
        self.text_encoder = nn.Sequential(
            nn.Linear(text_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
        )
        
        # Audio encoder
        self.audio_encoder = nn.Sequential(
            nn.Linear(audio_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
        )
        
        # Fusion layer
        self.fusion = nn.Sequential(
            nn.Linear(hidden_dim + hidden_dim // 2, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
        )
        
        # Classification head
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, num_classes),
        )
    
    def forward(
        self,
        text_features: torch.Tensor,
        audio_features: torch.Tensor,
    ) -> torch.Tensor:
        """
        Forward pass.
        
        Args:
            text_features: Text embeddings (batch, text_dim)
            audio_features: Audio features (batch, audio_dim)
        
        Returns:
            Logits (batch, num_classes)
        """
        # Encode each modality
        text_encoded = self.text_encoder(text_features)
        audio_encoded = self.audio_encoder(audio_features)
        
        # Concatenate and fuse
        fused = torch.cat([text_encoded, audio_encoded], dim=-1)
        fused = self.fusion(fused)
        
        # Classify
        logits = self.classifier(fused)
        
        return logits
    
    def predict(
        self,
        text_features: torch.Tensor,
        audio_features: Optional[torch.Tensor] = None,
    ) -> Dict[str, Any]:
        """Make predictions."""
        self.eval()
        
        if audio_features is None:
            # Text-only mode
            audio_features = torch.zeros(
                text_features.size(0), self.audio_dim
            ).to(text_features.device)
        
        with torch.no_grad():
            logits = self.forward(text_features, audio_features)
            probs = torch.softmax(logits, dim=-1)
        
        sentiment_labels = ["negative", "neutral", "positive", "sarcastic"]
        
        predictions = []
        for i, probs_i in enumerate(probs):
            pred_idx = probs_i.argmax().item()
            predictions.append({
                "sentiment": sentiment_labels[pred_idx],
                "confidence": probs_i[pred_idx].item(),
                "probabilities": {
                    label: probs_i[j].item()
                    for j, label in enumerate(sentiment_labels)
                },
            })
        
        return {"predictions": predictions}


class CrossModalAttention(nn.Module):
    """Cross-modal attention for audio-text fusion."""
    
    def __init__(
        self,
        query_dim: int,
        key_dim: int,
        hidden_dim: int,
        num_heads: int = 4,
    ):
        super().__init__()
        
        self.num_heads = num_heads
        self.head_dim = hidden_dim // num_heads
        
        self.query = nn.Linear(query_dim, hidden_dim)
        self.key = nn.Linear(key_dim, hidden_dim)
        self.value = nn.Linear(key_dim, hidden_dim)
        
        self.output = nn.Linear(hidden_dim, hidden_dim)
        self.scale = self.head_dim ** -0.5
    
    def forward(
        self,
        query: torch.Tensor,
        key_value: torch.Tensor,
    ) -> torch.Tensor:
        """Cross-attention forward pass."""
        batch_size = query.size(0)
        
        # Linear projections
        Q = self.query(query).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.key(key_value).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.value(key_value).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        
        # Attention scores
        scores = torch.matmul(Q, K.transpose(-2, -1)) * self.scale
        attention = torch.softmax(scores, dim=-1)
        
        # Apply attention to values
        context = torch.matmul(attention, V)
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.head_dim)
        
        return self.output(context)


if __name__ == "__main__":
    print("Testing MultiModalSentimentModel...")
    
    model = MultiModalSentimentModel(
        text_dim=768,
        audio_dim=8,
        hidden_dim=256,
        num_classes=4,
    )
    
    # Mock inputs
    text_features = torch.randn(2, 768)
    audio_features = torch.randn(2, 8)
    
    logits = model(text_features, audio_features)
    print(f"Output shape: {logits.shape}")
    print(f"Total parameters: {model.get_num_parameters():,}")