import torch
import torch.nn as nn
import torchaudio
from typing import Tuple

class DriverSpeechEmotionDetector(nn.Module):
    """
    驾驶员语音情感检测器
    
    论文: IEEE Trans. Intelligent Vehicles
    架构: 全局(CNN) + 局部(LSTM) 特征融合
    
    情绪类别: neutral, happy, sad, angry, fear, surprise
    """
    
    def __init__(self, num_emotions=6, sample_rate=16000):
        super().__init__()
        self.sample_rate = sample_rate
        
        # 音频预处理
        self.mel_transform = torchaudio.transforms.MelSpectrogram(
            sample_rate=sample_rate,
            n_fft=1024,
            hop_length=512,
            n_mels=80
        )
        
        # 全局分支：CNN处理频谱图
        self.global_branch = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),
            
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),
            
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1))
        )
        
        # 局部分支：LSTM处理时序特征
        self.local_branch = nn.LSTM(
            input_size=80,  # mel bands
            hidden_size=128,
            num_layers=2,
            batch_first=True,
            bidirectional=True,
            dropout=0.3
        )
        
        # 融合+分类
        self.classifier = nn.Sequential(
            nn.Linear(128 + 256, 256),  # global + local
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, num_emotions)
        )
        
        # 压力等级回归（0-1）
        self.stress_regressor = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )
        
    def forward(self, audio: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            audio: 音频波形, shape=(B, 1, samples)
            
        Returns:
            emotion_logits: (B, num_emotions)
            stress_level: (B, 1) 0-1
        """
        # 提取Mel频谱图
        mel = self.mel_transform(audio.squeeze(1))  # (B, n_mels, time)
        mel = torch.log(mel + 1e-9)  # log-mel
        
        # 全局特征
        mel_2d = mel.unsqueeze(1)  # (B, 1, n_mels, time)
        global_feat = self.global_branch(mel_2d)  # (B, 128, 1, 1)
        global_feat = global_feat.flatten(1)  # (B, 128)
        
        # 局部特征
        mel_t = mel.transpose(1, 2)  # (B, time, n_mels)
        local_out, _ = self.local_branch(mel_t)  # (B, time, 256)
        local_feat = local_out[:, -1, :]  # 取最后时刻 (B, 256)
        
        # 融合
        fused = torch.cat([global_feat, local_feat], dim=1)  # (B, 384)
        
        # 分类
        emotion_logits = self.classifier(fused)
        stress_level = self.stress_regressor(fused[:, :128])
        
        return emotion_logits, stress_level


class CabinVoicePipeline:
    """
    车载语音情感检测完整管道
    
    挑战:
    1. 车内噪声（发动机/风/路面）
    2. 远场采集（1-2米距离）
    3. EV座舱声学不同于燃油车
    4. 多乘员干扰
    """
    
    # 情绪→IMS状态映射
    EMOTION_TO_STATE = {
        'neutral':  {'stress': 0.2, 'alert': True,  'action': 'none'},
        'happy':    {'stress': 0.1, 'alert': True,  'action': 'none'},
        'sad':      {'stress': 0.4, 'alert': False, 'action': 'monitor'},
        'angry':    {'stress': 0.8, 'alert': True,  'action': 'warn'},
        'fear':     {'stress': 0.9, 'alert': True,  'action': 'alert'},
        'surprise': {'stress': 0.6, 'alert': True,  'action': 'monitor'},
    }
    
    def __init__(self):
        self.model = DriverSpeechEmotionDetector()
        self.model.eval()
        
        # 车内噪声阈值
        self.SNR_THRESHOLD = 5  # 最低5dB信噪比
        
        # EV特有问题：低频路面噪声与男性语音重叠
        self.ROAD_NOISE_BAND = (100, 500)  # Hz
        
    def process_utterance(self, audio: np.ndarray, 
                          speaker_location: str = 'driver') -> dict:
        """
        处理一段语音
        
        Args:
            audio: 音频波形
            speaker_location: 'driver'/'passenger'/'rear'
            
        Returns:
            dict: 情绪+压力+IMS建议
        """
        if speaker_location != 'driver':
            return {'skip': True, 'reason': 'non_driver_speech'}
        
        # SNR检查
        snr = self._compute_snr(audio)
        if snr < self.SNR_THRESHOLD:
            return {'skip': True, 'reason': f'low_snr_{snr:.1f}dB'}
        
        # 推理
        with torch.no_grad():
            audio_t = torch.FloatTensor(audio).unsqueeze(0).unsqueeze(0)
            emotion_logits, stress = self.model(audio_t)
            
            emotion_probs = torch.softmax(emotion_logits, dim=1)
            emotion_idx = emotion_probs.argmax(1).item()
            emotions = ['neutral', 'happy', 'sad', 'angry', 'fear', 'surprise']
            emotion = emotions[emotion_idx]
            
            state = self.EMOTION_TO_STATE[emotion]
            
            return {
                'emotion': emotion,
                'emotion_confidence': emotion_probs[0, emotion_idx].item(),
                'stress_level': stress.item(),
                'ims_action': state['action'],
                'alert': state['alert'],
                'snr_db': snr
            }
    
    def _compute_snr(self, audio: np.ndarray) -> float:
        """计算信噪比"""
        # 简化版：假设前100ms为噪声
        if len(audio) < 1600:
            return 0
        noise = audio[:1600]
        signal = audio[1600:]
        
        noise_power = np.mean(noise**2)
        signal_power = np.mean(signal**2)
        
        if noise_power == 0:
            return 100
        
        return 10 * np.log10(signal_power / noise_power)

# Euro NCAP 2026 语音相关要求
ENCAP_VOICE_REQUIREMENTS = {
    'touchscreen_penalty': '核心功能仅触屏→扣分',
    'voice_accepted': '语音是接受的替代输入',
    'recognition_quality': '识别质量=安全话题',
    'implementation': '语音命令不得需要驾驶员低头看屏幕',
}

# EU AI Act 对情感识别的规定
EU_AI_ACT = {
    'emotion_recognition_ban': '工作场所和教育场景禁止',
    'safety_exception': '疲劳检测用于事故预防不在此定义内',
    'recital_18': '检测疲劳等物理状态→事故预防→豁免',
    'cabin_systems': '需项目级法律分类',
    'requirements': ['知情同意', '透明度', '数据最小化', 'GDPR合规']
}

class MultimodalDriverState:
    """
    视觉+语音+生理多模态驾驶员状态
    
    语音补充视觉无法检测的情绪维度
    """
    
    def __init__(self):
        self.weights = {
            'visual_fatigue': 0.35,   # PERCLOS
            'visual_distraction': 0.20, # 视线偏离
            'voice_emotion': 0.15,    # 语音情绪
            'voice_stress': 0.15,     # 语音压力
            'hrv': 0.15               # 心率变异性
        }
        
    def assess(self, visual_fatigue, visual_distraction,
               voice_emotion, voice_stress, hrv):
        """
        多模态评估
        
        Returns:
            state: 'normal'|'stressed'|'fatigued'|'dangerous'
        """
        scores = {
            'visual_fatigue': visual_fatigue,
            'visual_distraction': visual_distraction,
            'voice_emotion': voice_emotion,
            'voice_stress': voice_stress,
            'hrv': hrv
        }
        
        total = sum(self.weights[k] * scores[k] for k in scores)
        
        if total > 0.7:
            return 'dangerous'
        elif total > 0.5:
            return 'stressed'
        elif total > 0.3:
            return 'fatigued'
        else:
            return 'normal'