import torch
import torch.nn as nn
from typing import Dict, Optional

class InCarEmoBaseline(nn.Module):
    """
    InCarEmo 多模态基线模型
    
    支持:
    - 单模态: RGB / IR / Audio / Text
    - 多模态融合: 任意组合
    - 模态缺失测试
    - 噪声鲁棒性测试
    """
    
    def __init__(
        self,
        modalities: list = ['rgb', 'ir', 'audio', 'text'],
        hidden_dim: int = 256,
        num_emotions: int = 4,
        num_fatigue: int = 3,
        num_distraction: int = 3
    ):
        super().__init__()
        self.modalities = modalities
        
        # 各模态编码器
        self.encoders = nn.ModuleDict()
        if 'rgb' in modalities:
            self.encoders['rgb'] = VideoEncoder(input_channels=3, hidden_dim=hidden_dim)
        if 'ir' in modalities:
            self.encoders['ir'] = VideoEncoder(input_channels=1, hidden_dim=hidden_dim)
        if 'audio' in modalities:
            self.encoders['audio'] = AudioEncoder(hidden_dim=hidden_dim)
        if 'text' in modalities:
            self.encoders['text'] = TextEncoder(hidden_dim=hidden_dim)
        
        # 跨模态注意力
        self.cross_attn = CrossModalAttention(
            hidden_dim=hidden_dim,
            num_heads=4
        )
        
        # 多任务头
        self.emotion_head = nn.Linear(hidden_dim, num_emotions)
        self.fatigue_head = nn.Linear(hidden_dim, num_fatigue)
        self.distraction_head = nn.Linear(hidden_dim, num_distraction)
    
    def forward(
        self,
        rgb: Optional[torch.Tensor] = None,
        ir: Optional[torch.Tensor] = None,
        audio: Optional[torch.Tensor] = None,
        text: Optional[torch.Tensor] = None
    ) -> Dict[str, torch.Tensor]:
        """前向传播，支持模态缺失"""
        inputs = {'rgb': rgb, 'ir': ir, 'audio': audio, 'text': text}
        
        # 编码可用模态
        embeddings = []
        for mod in self.modalities:
            if inputs[mod] is not None:
                emb = self.encoders[mod](inputs[mod])
                embeddings.append(emb)
        
        if not embeddings:
            raise ValueError("至少需要一个模态输入")
        
        # 跨模态融合
        fused = self.cross_attn(embeddings)
        
        # 多任务输出
        return {
            'emotion': self.emotion_head(fused),
            'fatigue': self.fatigue_head(fused),
            'distraction': self.distraction_head(fused)
        }


class VideoEncoder(nn.Module):
    def __init__(self, input_channels, hidden_dim):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(input_channels, 32, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(64, hidden_dim)
        )
    def forward(self, x):
        return self.conv(x)


class AudioEncoder(nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(128, hidden_dim),  # 梅尔频谱输入
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )
    def forward(self, x):
        return self.encoder(x)


class TextEncoder(nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(768, hidden_dim),  # BERT嵌入输入
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )
    def forward(self, x):
        return self.encoder(x)


class CrossModalAttention(nn.Module):
    def __init__(self, hidden_dim, num_heads):
        super().__init__()
        self.attn = nn.MultiheadAttention(hidden_dim, num_heads, batch_first=True)
        self.norm = nn.LayerNorm(hidden_dim)
    
    def forward(self, embeddings: list) -> torch.Tensor:
        # 将各模态嵌入堆叠为序列
        seq = torch.stack(embeddings, dim=1)  # (B, M, D)
        attended, _ = self.attn(seq, seq, seq)
        pooled = self.norm(attended.mean(dim=1))
        return pooled


# ============ 模态缺失测试 ============

def test_modality_missing():
    """测试模态缺失场景下的性能"""
    model = InCarEmoBaseline(modalities=['rgb', 'ir', 'audio', 'text'])
    
    # 完整模态
    full_input = {
        'rgb': torch.randn(4, 3, 224, 224),
        'ir': torch.randn(4, 1, 224, 224),
        'audio': torch.randn(4, 128),
        'text': torch.randn(4, 768)
    }
    
    # 缺失音频
    no_audio = {**full_input, 'audio': None}
    
    # 只有RGB
    rgb_only = {'rgb': full_input['rgb'], 'ir': None, 'audio': None, 'text': None}
    
    # 夜间（无RGB，有IR）
    night = {'rgb': None, 'ir': full_input['ir'], 'audio': full_input['audio'], 'text': None}
    
    results = {}
    for name, data in [('full', full_input), ('no_audio', no_audio), 
                        ('rgb_only', rgb_only), ('night', night)]:
        out = model(**data)
        results[name] = {
            'emotion': out['emotion'].shape,
            'fatigue': out['fatigue'].shape,
            'distraction': out['distraction'].shape
        }
    
    return results


if __name__ == "__main__":
    results = test_modality_missing()
    for name, shapes in results.items():
        print(f"{name}: {shapes}")
    
    print("\n=== InCarEmo 数据集关键指标 ===")
    print("模态: RGB + IR + Audio + Text（首个四模态车内数据集）")
    print("任务: 情感识别 + 疲劳检测 + 分心监控")
    print("情绪类别: 愤怒/悲伤/激动/喜悦")
    print("语言: 中文（主）+ 英文（辅助）")
    print("光照: 白天/夜间/隧道/阴影")
    print("特色: 模态缺失测试 + 噪声鲁棒性评估")