import torch
import torch.nn as nn

class InCarEmoFusionModel(nn.Module):
    """
    InCarEmo 多模态融合模型
    
    融合4个模态：面部+语音+生理+驾驶行为
    """
    
    def __init__(self, n_emotions=7):
        super().__init__()
        
        # 1. 面部分支（CNN+LSTM）
        self.face_encoder = nn.Sequential(
            nn.Conv2d(3, 32, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.ReLU(),
            nn.AdaptiveAvgPool2d((4, 4)),
        )
        self.face_lstm = nn.LSTM(128*16, 256, batch_first=True)
        
        # 2. 语音分支（CNN）
        self.voice_encoder = nn.Sequential(
            nn.Conv1d(80, 128, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv1d(128, 256, 3, stride=2, padding=1), nn.ReLU(),
            nn.AdaptiveAvgPool1d(1),
        )
        
        # 3. 生理分支（PPG+EDA）
        self.phys_encoder = nn.Sequential(
            nn.Linear(200, 128), nn.ReLU(),  # 2s窗口@100Hz
            nn.Linear(128, 64), nn.ReLU(),
        )
        
        # 4. 驾驶行为分支
        self.drive_encoder = nn.Sequential(
            nn.Linear(120, 64), nn.ReLU(),  # 2s窗口@60Hz × 1维
            nn.Linear(64, 32), nn.ReLU(),
        )
        
        # 融合层
        self.fusion = nn.Sequential(
            nn.Linear(256 + 256 + 64 + 32, 256), nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, n_emotions)
        )
        
        # 效价/唤醒度回归头
        self.va_head = nn.Linear(256, 2)  # (valence, arousal)
    
    def forward(self, face, voice, phys, driving):
        # 面部
        b, t, c, h, w = face.shape
        face_feat = self.face_encoder(face.view(b*t, c, h, w))
        face_feat = face_feat.view(b, t, -1)
        _, (hn, _) = self.face_lstm(face_feat)
        face_out = hn[-1]  # (B, 256)
        
        # 语音
        voice_out = self.voice_encoder(voice).squeeze(-1)  # (B, 256)
        
        # 生理
        phys_out = self.phys_encoder(phys)  # (B, 64)
        
        # 驾驶
        drive_out = self.drive_encoder(driving)  # (B, 32)
        
        # 融合
        combined = torch.cat([face_out, voice_out, phys_out, drive_out], dim=1)
        emotions = self.fusion(combined)
        va = self.va_head(combined)
        
        return {
            'emotions': emotions,  # (B, 7) 分类
            'valence_arousal': va,  # (B, 2) 回归
        }


# 测试
if __name__ == "__main__":
    model = InCarEmoFusionModel()
    
    # 模拟输入
    face = torch.randn(2, 30, 3, 64, 64)    # 30帧面部视频
    voice = torch.randn(2, 80, 160)         # 80维MFCC, 1s
    phys = torch.randn(2, 200)              # 2s生理信号
    driving = torch.randn(2, 120)           # 2s驾驶行为
    
    output = model(face, voice, phys, driving)
    print(f"情绪分类: {output['emotions'].shape}")
    print(f"效价/唤醒度: {output['valence_arousal'].shape}")
    print(f"参数总数: {sum(p.numel() for p in model.parameters()):,}")