import numpy as np
import librosa
from typing import Tuple, Dict

class SpeechEmotionRecognizer:
    """
    车载语音情感识别系统
    
    综述: ScienceDirect 2026 系统综述
    特征: 384维 (eGeMAPS 88 + 频谱 128 + 深度 168)
    模型: CNN + Transformer hybrid
    
    应用: 压力检测 + 情绪分类 → DMS/OMS融合
    """
    
    # 7类情绪（Ford专利定义）
    EMOTIONS = ['anger', 'fear', 'sadness', 'happiness', 
                'surprise', 'disgust', 'neutral']
    
    # 压力等级映射
    STRESS_MAP = {
        'anger': 0.9,      # 高压力
        'fear': 0.85,     # 高压力
        'sadness': 0.5,   # 中压力
        'disgust': 0.6,   # 中压力
        'surprise': 0.4,  # 低压力
        'happiness': 0.1, # 低压力
        'neutral': 0.2    # 基线
    }
    
    def __init__(self, sr=16000):
        self.sr = sr
        
    def extract_features(self, audio: np.ndarray) -> Dict[str, np.ndarray]:
        """
        提取384维语音特征
        
        三类特征:
        1. eGeMAPS声学特征 (88维) - 专业语音分析
        2. 频谱特征 (128维) - MFCC + delta
        3. 深度特征 (168维) - wav2vec2 embeddings
        
        Args:
            audio: 语音信号, shape=(N,)
            
        Returns:
            dict: 三类特征
        """
        # 1. 声学特征（eGeMAPS子集）
        acoustic = self._extract_acoustic(audio)
        
        # 2. 频谱特征
        spectral = self._extract_spectral(audio)
        
        # 3. 深度特征（简化版，实际需wav2vec2）
        deep = self._extract_deep_features(audio)
        
        return {
            'acoustic': acoustic,      # (88,)
            'spectral': spectral,      # (128,)
            'deep': deep,              # (168,)
            'combined': np.concatenate([
                acoustic, spectral, deep
            ])  # (384,)
        }
    
    def _extract_acoustic(self, audio: np.ndarray) -> np.ndarray:
        """提取eGeMAPS声学特征子集"""
        features = []
        
        # F0基频统计
        f0, voiced_flag, _ = librosa.pyin(
            audio, fmin=80, fmax=400, sr=self.sr
        )
        f0_valid = f0[~np.isnan(f0)]
        features.extend([
            np.mean(f0_valid) if len(f0_valid) > 0 else 0,
            np.std(f0_valid) if len(f0_valid) > 0 else 0,
            np.percentile(f0_valid, 20) if len(f0_valid) > 0 else 0,
            np.percentile(f0_valid, 80) if len(f0_valid) > 0 else 0,
        ])
        
        # 能量特征
        rms = librosa.feature.rms(y=audio)[0]
        features.extend([
            np.mean(rms), np.std(rms),
            np.percentile(rms, 20), np.percentile(rms, 80)
        ])
        
        # Jitter（频率微扰）
        if len(f0_valid) > 1:
            jitter = np.mean(np.abs(np.diff(f0_valid)) / f0_valid[:-1])
        else:
            jitter = 0
        features.append(jitter)
        
        # Shimmer（幅度微扰）
        if len(rms) > 1:
            shimmer = np.mean(np.abs(np.diff(rms)) / rms[:-1])
        else:
            shimmer = 0
        features.append(shimmer)
        
        # 谐波噪声比
        harmonic, percussive = librosa.effects.hpss(audio)
        hnr = np.mean(np.abs(harmonic)) / (np.mean(np.abs(percussive)) + 1e-10)
        features.append(hnr)
        
        # 补齐到88维
        while len(features) < 88:
            features.append(0)
            
        return np.array(features[:88])
    
    def _extract_spectral(self, audio: np.ndarray) -> np.ndarray:
        """提取频谱特征（MFCC + delta）"""
        # MFCC
        mfcc = librosa.feature.mfcc(y=audio, sr=self.sr, n_mfcc=40)
        
        # Delta (一阶差分)
        delta = librosa.feature.delta(mfcc)
        
        # Delta-delta (二阶差分)
        delta2 = librosa.feature.delta(mfcc, order=2)
        
        # 统计特征
        features = []
        for feat in [mfcc, delta, delta2]:
            features.extend(np.mean(feat, axis=1))
            features.extend(np.std(feat, axis=1))
        
        # 补齐/截断到128维
        features = features[:128]
        while len(features) < 128:
            features.append(0)
            
        return np.array(features)
    
    def _extract_deep_features(self, audio: np.ndarray) -> np.ndarray:
        """
        深度特征提取（简化版）
        
        实际应用需 wav2vec2 预训练模型
        此处用自编码器特征替代
        """
        # 短时傅里叶变换
        stft = np.abs(librosa.stft(audio, n_fft=512, hop_length=160))
        
        # Mel频谱
        mel = librosa.feature.melspectrogram(
            S=stft**2, sr=self.sr, n_mels=64
        )
        mel_db = librosa.power_to_db(mel)
        
        # 统计特征
        features = np.concatenate([
            np.mean(mel_db, axis=1),
            np.std(mel_db, axis=1),
            np.percentile(mel_db, 20, axis=1),
        ])  # 64*3 = 192 -> 截断到168
        
        return features[:168]
    
    def classify_emotion(self, audio: np.ndarray) -> Dict:
        """
        情绪分类
        
        Returns:
            dict: 情绪概率 + 压力评分
        """
        features = self.extract_features(audio)
        
        # 模拟分类器输出（实际应用需训练模型）
        # 基于声学特征的简单规则
        f0_var = features['acoustic'][1]  # F0标准差
        energy = features['acoustic'][4]  # RMS均值
        jitter = features['acoustic'][8]   # Jitter
        
        # 简单规则分类（实际需深度学习模型）
        if f0_var > 50 and energy > 0.05:
            emotion_probs = {'anger': 0.6, 'fear': 0.2, 'neutral': 0.2}
        elif f0_var < 10 and energy < 0.02:
            emotion_probs = {'sadness': 0.5, 'neutral': 0.4, 'happiness': 0.1}
        else:
            emotion_probs = {'neutral': 0.6, 'happiness': 0.2, 'surprise': 0.2}
        
        # 归一化
        total = sum(emotion_probs.values())
        emotion_probs = {k: v/total for k, v in emotion_probs.items()}
        
        # 压力评分
        stress_score = sum(
            self.STRESS_MAP.get(emo, 0.3) * prob 
            for emo, prob in emotion_probs.items()
        )
        
        return {
            'emotion_probs': emotion_probs,
            'dominant_emotion': max(emotion_probs, key=emotion_probs.get),
            'stress_score': stress_score,
            'stress_level': 'high' if stress_score > 0.6 else 
                          'moderate' if stress_score > 0.3 else 'low',
            'features_dim': len(features['combined'])
        }


# 测试
if __name__ == "__main__":
    recognizer = SpeechEmotionRecognizer(sr=16000)
    
    # 模拟3秒语音
    duration = 3.0
    t = np.linspace(0, duration, int(16000 * duration))
    
    # 正常语音（中等F0变化）
    normal_audio = (0.3 * np.sin(2*np.pi*120*t) +     # 基频120Hz
                    0.1 * np.sin(2*np.pi*240*t) +      # 谐波
                    0.02 * np.random.randn(len(t)))     # 噪声
    
    # 愤怒语音（高F0变化+高能量）
    angry_audio = (0.8 * np.sin(2*np.pi*180*t * (1 + 0.1*np.sin(2*np.pi*5*t))) +
                   0.3 * np.sin(2*np.pi*360*t) +
                   0.05 * np.random.randn(len(t)))
    
    print("=== 正常语音 ===")
    result_normal = recognizer.classify_emotion(normal_audio)
    print(f"主导情绪: {result_normal['dominant_emotion']}")
    print(f"压力评分: {result_normal['stress_score']:.2f}")
    print(f"压力等级: {result_normal['stress_level']}")
    
    print("\n=== 愤怒语音 ===")
    result_angry = recognizer.classify_emotion(angry_audio)
    print(f"主导情绪: {result_angry['dominant_emotion']}")
    print(f"压力评分: {result_angry['stress_score']:.2f}")
    print(f"压力等级: {result_angry['stress_level']}")

class MultimodalEmotionFusion:
    """
    多模态情绪融合：视觉+语音+生理
    
    权重动态调整：
    - 正常光照: 视觉40% + 语音30% + 生理30%
    - 低光/遮挡: 语音50% + 生理30% + 视觉20%
    - 语音静默: 视觉50% + 生理30% + 行为20%
    """
    
    DEFAULT_WEIGHTS = {
        'visual': 0.40,
        'voice': 0.30,
        'physio': 0.30
    }
    
    def __init__(self):
        self.weights = self.DEFAULT_WEIGHTS.copy()
        self.signal_quality = {k: 1.0 for k in self.DEFAULT_WEIGHTS}
        
    def update_weights(self, lighting_ok: bool, voice_active: bool, 
                       sensors_ok: bool):
        """根据信号质量动态调整权重"""
        if not lighting_ok:
            self.signal_quality['visual'] = 0.3
            self.weights['voice'] = 0.50
            self.weights['physio'] = 0.20
        else:
            self.signal_quality['visual'] = 1.0
            
        if not voice_active:
            self.signal_quality['voice'] = 0.0
            self.weights['visual'] = 0.55
            self.weights['physio'] = 0.35
            
        if not sensors_ok:
            self.signal_quality['physio'] = 0.0
            
        # 归一化
        total = sum(self.weights.values())
        if total > 0:
            self.weights = {k: v/total for k, v in self.weights.items()}
            
    def fuse(self, visual_score: float, voice_score: float, 
             physio_score: float) -> dict:
        """加权融合"""
        fused = (visual_score * self.weights['visual'] +
                voice_score * self.weights['voice'] +
                physio_score * self.weights['physio'])
        
        return {
            'stress_score': fused,
            'weights': self.weights,
            'quality': self.signal_quality,
            'level': 'high' if fused > 0.6 else 
                    'moderate' if fused > 0.3 else 'low'
        }