import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class MHLTDriverEmotionModel(nn.Module):
    """
    MHLT: Multi-Head attention Level fusion for Driver emotion recognition
    
    论文: Frontiers in Physics, 2025
    核心: 视频(面部表情) + 音频(语音情绪) → 模型级融合 → 情绪强度回归
    
    创新点:
    1. 情绪连续维度 [-1, +1] 而非离散分类
    2. 多头注意力模型级融合
    3. 同时预测情绪类型和强度
    """
    
    def __init__(self, 
                 video_feature_dim=512, 
                 audio_feature_dim=256,
                 hidden_dim=256,
                 num_heads=8):
        super().__init__()
        
        # 视频编码器 (基于VGG/ResNet的面部表情特征)
        self.video_encoder = nn.Sequential(
            nn.Linear(video_feature_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU()
        )
        
        # 音频编码器 (基于Transformer的语音情绪特征)
        self.audio_encoder = nn.Sequential(
            nn.Linear(audio_feature_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU()
        )
        
        # 多头注意力融合层
        self.fusion_attention = nn.MultiheadAttention(
            embed_dim=hidden_dim * 2,
            num_heads=num_heads,
            dropout=0.1,
            batch_first=True
        )
        
        # 情绪分类头 (正/中/负)
        self.classification_head = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.GELU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, 3)  # 3类: 正/中/负
        )
        
        # 情绪强度回归头 [-1, +1]
        self.intensity_head = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.GELU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, 1),
            nn.Tanh()  # 输出 [-1, 1]
        )
        
    def forward(self, video_features: torch.Tensor, 
                audio_features: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        前向传播
        
        Args:
            video_features: 面部表情特征 [B, T, 512]
            audio_features: 语音情绪特征 [B, T, 256]
            
        Returns:
            emotion_class: 情绪分类logits [B, 3]
            emotion_intensity: 情绪强度 [-1, 1] [B, 1]
        """
        # 编码
        v_encoded = self.video_encoder(video_features)  # [B, T, 256]
        a_encoded = self.audio_encoder(audio_features)  # [B, T, 256]
        
        # 拼接
        fused = torch.cat([v_encoded, a_encoded], dim=-1)  # [B, T, 512]
        
        # 多头注意力融合
        attended, _ = self.fusion_attention(
            fused, fused, fused
        )
        
        # 池化
        pooled = attended.mean(dim=1)  # [B, 512]
        
        # 双头输出
        emotion_class = self.classification_head(pooled)
        emotion_intensity = self.intensity_head(pooled)
        
        return emotion_class, emotion_intensity


# 训练和测试代码
class MHLTTrainer:
    """MHLT模型训练器"""
    
    def __init__(self, model, lr=1e-4):
        self.model = model
        self.optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
        
        # 分类损失
        self.cls_criterion = nn.CrossEntropyLoss()
        # 强度回归损失 (MSE)
        self.reg_criterion = nn.MSELoss()
        
    def train_step(self, video_feat, audio_feat, emotion_label, intensity_label):
        """
        单步训练
        
        Args:
            video_feat: [B, T, 512]
            audio_feat: [B, T, 256]
            emotion_label: [B] (0=正, 1=中, 2=负)
            intensity_label: [B] (-1.0 to 1.0)
        """
        self.optimizer.zero_grad()
        
        cls_pred, int_pred = self.model(video_feat, audio_feat)
        
        # 总损失 = 分类损失 + 强度损失
        cls_loss = self.cls_criterion(cls_pred, emotion_label)
        reg_loss = self.reg_criterion(int_pred.squeeze(), intensity_label)
        
        # 论文使用加权组合
        total_loss = cls_loss + 0.5 * reg_loss
        
        total_loss.backward()
        self.optimizer.step()
        
        return {
            'total_loss': total_loss.item(),
            'cls_loss': cls_loss.item(),
            'reg_loss': reg_loss.item()
        }


if __name__ == "__main__":
    # 模型初始化
    model = MHLTDriverEmotionModel(
        video_feature_dim=512,
        audio_feature_dim=256,
        hidden_dim=256,
        num_heads=8
    )
    
    # 模拟输入
    batch_size = 4
    seq_len = 30  # 30帧
    
    video_feat = torch.randn(batch_size, seq_len, 512)
    audio_feat = torch.randn(batch_size, seq_len, 256)
    
    # 前向传播
    cls_pred, int_pred = model(video_feat, audio_feat)
    
    print("=== MHLT模型测试 ===")
    print(f"输入: 视频{video_feat.shape} + 音频{audio_feat.shape}")
    print(f"分类输出: {cls_pred.shape} (3类: 正/中/负)")
    print(f"强度输出: {int_pred.shape} (范围: [-1, +1])")
    print(f"强度值: {int_pred.detach().numpy().flatten()}")
    
    # 训练模拟
    trainer = MHLTTrainer(model)
    labels = torch.tensor([2, 1, 0, 2])  # 负, 中, 正, 负
    intensities = torch.tensor([-0.7, 0.0, 0.3, -0.8])
    
    for epoch in range(5):
        result = trainer.train_step(video_feat, audio_feat, labels, intensities)
        print(f"Epoch {epoch+1}: loss={result['total_loss']:.4f} "
              f"(cls={result['cls_loss']:.4f}, reg={result['reg_loss']:.4f})")
    
    # 情绪强度解释
    print("\n=== 情绪强度解读 ===")
    print("-0.8: 极端负面 → 可能危险驾驶（愤怒/恐慌）")
    print("-0.2: 轻度负面 → 可能更谨慎驾驶")
    print(" 0.0: 中性 → 正常驾驶状态")
    print("+0.3: 轻度正面 → 放松但可能分心")
    print("+0.8: 极端正面 → 兴奋可能导致冒险")

class DMSEmotionIntegration:
    """
    DMS + 情绪识别集成方案
    
    当情绪强度超出安全范围时触发干预
    """
    
    SAFE_RANGE = (-0.2, 0.2)
    WARNING_RANGE = (-0.6, -0.2, 0.2, 0.6)
    DANGER_RANGE = (-1.0, -0.6, 0.6, 1.0)
    
    def __init__(self):
        self.emotion_model = MHLTDriverEmotionModel()
        self.current_emotion = 0.0
        self.emotion_history = []
        
    def monitor(self, video_feat, audio_feat) -> dict:
        """实时情绪监控"""
        with torch.no_grad():
            cls_pred, int_pred = self.emotion_model(video_feat, audio_feat)
            
        intensity = int_pred.item()
        emotion_class = ['positive', 'neutral', 'negative'][cls_pred.argmax().item()]
        
        self.emotion_history.append(intensity)
        
        # 滑动窗口平均（10秒）
        if len(self.emotion_history) > 300:
            recent = self.emotion_history[-300:]
            avg_intensity = sum(recent) / len(recent)
        else:
            avg_intensity = intensity
        
        return self._generate_response(avg_intensity, emotion_class)
    
    def _generate_response(self, intensity: float, emotion: str) -> dict:
        """生成情绪干预响应"""
        if intensity < -0.6:
            return {
                'level': 'danger',
                'action': '强烈建议靠边停车',
                'audio': '舒缓音乐',
                'hud': '请冷静，建议休息'
            }
        elif intensity < -0.2:
            return {
                'level': 'warning',
                'action': '提醒注意情绪',
                'audio': '轻音乐',
                'hud': '保持冷静驾驶'
            }
        elif intensity > 0.6:
            return {
                'level': 'warning',
                'action': '提醒注意分心',
                'audio': None,
                'hud': '请专注驾驶'
            }
        else:
            return {'level': 'safe', 'action': None, 'audio': None, 'hud': None}