"""
座舱边缘语音 AI 代理管道
从麦克风阵列到语音响应的完整流程

参考：
1. Mercedes-Benz × Liquid AI: MB.OS 集成（2026 H2）
2. Hyundai Gleo AI: 全新 Grandeur 量产
3. Lucid × SoundHound: 离线对话
4. ENERZAi: 超轻量化语音AI

关键指标：
- STT延迟：<300ms（边缘）vs 1-5s（云端）
- LLM推理：<200ms（量化模型）
- TTS延迟：<100ms（流式合成）
- 总延迟：<500ms（用户感知自然）
- 内存占用：<2GB（车规AP约束）
- 离线运行：100%功能可用
"""

import numpy as np
from dataclasses import dataclass
from typing import Optional, List

@dataclass
class VoiceAgentConfig:
    """座舱语音代理配置"""
    # 音频参数
    sample_rate: int = 16000
    frame_length: int = 512  # 32ms at 16kHz
    hop_length: int = 160    # 10ms hop
    
    # 麦克风阵列
    num_mics: int = 4  # 4麦线性阵列
    mic_spacing: float = 0.04  # 4cm间距
    
    # STT 参数
    vocab_size: int = 5000  # 限制词表（车规场景）
    max_utterance_sec: int = 10
    
    # 情绪识别
    num_emotions: int = 7  # neutral/happy/sad/angry/fear/surprise/disgust
    
    # 部署约束
    max_memory_mb: int = 2048  # 2GB
    max_latency_ms: int = 500  # 总延迟
    
    # 离线运行
    offline_mode: bool = True


class CabinVoiceAgent:
    """
    座舱语音 AI 代理
    
    完整管道：
    1. 多麦克风波束形成 → 声源定位+降噪
    2. VAD → 检测语音段
    3. STT → 语音转文字（边缘量化模型）
    4. 情绪识别 → 声学情绪分类
    5. NLU/LLM → 意图理解
    6. DM → 对话管理
    7. TTS → 语音合成
    
    座舱特殊处理：
    - 区域识别：判断谁在说话（驾驶员/副驾/后排）
    - 噪声环境：发动机+风噪+路噪
    - 多语言：中文+英文+方言
    - 上下文感知：驾驶状态+车辆状态
    """
    
    def __init__(self, config: VoiceAgentConfig):
        self.config = config
        
    def beamforming(self, multi_mic_signal: np.ndarray,
                     look_direction: int = 0) -> np.ndarray:
        """
        波束形成：空间滤波+声源定位
        
        4麦线性阵列，延迟-叠加波束形成
        
        Args:
            multi_mic_signal: (num_mics, num_samples) 多麦克风信号
            look_direction: 目标方向（度），0=正前方
        
        Returns:
            enhanced_signal: (num_samples,) 增强后信号
        """
        num_mics, N = multi_mic_signal.shape
        c = 343.0  # 声速 m/s
        
        # 计算导向矢量
        d = self.config.mic_spacing
        angle_rad = np.deg2rad(look_direction)
        
        # 各麦克风时延
        delays = np.array([
            i * d * np.sin(angle_rad) / c 
            for i in range(num_mics)
        ])
        
        # 分数延迟滤波
        enhanced = np.zeros(N)
        for i in range(num_mics):
            delay_samples = delays[i] * self.config.sample_rate
            delay_int = int(np.round(delay_samples))
            delay_frac = delay_samples - delay_int
            
            # 线性插值分数延迟
            if delay_int >= 0:
                shifted = np.concatenate([
                    np.zeros(delay_int),
                    multi_mic_signal[i, :N-delay_int]
                ])
            else:
                shifted = np.concatenate([
                    multi_mic_signal[i, -delay_int:],
                    np.zeros(-delay_int)
                ])
            
            enhanced += shifted / num_mics
        
        return enhanced
    
    def voice_activity_detection(self, signal: np.ndarray) -> List[dict]:
        """
        语音活动检测（VAD）
        
        基于能量 + 零交叉率的双门限检测
        
        Returns:
            segments: [{start, end, duration}] 语音段列表
        """
        frame_size = self.config.frame_length
        hop = self.config.hop_length
        num_frames = (len(signal) - frame_size) // hop
        
        # 计算每帧能量
        energies = []
        zcrs = []
        for i in range(num_frames):
            frame = signal[i*hop : i*hop+frame_size]
            energy = np.sum(frame**2) / frame_size
            zcr = np.sum(np.diff(np.sign(frame)) != 0) / frame_size
            energies.append(energy)
            zcrs.append(zcr)
        
        energies = np.array(energies)
        zcrs = np.array(zcrs)
        
        # 自适应阈值
        energy_threshold = np.mean(energies) * 0.5
        zcr_threshold = np.mean(zcrs) * 1.5
        
        # 双门限检测
        voiced = (energies > energy_threshold) & (zcrs < zcr_threshold)
        
        # 连续段提取
        segments = []
        in_speech = False
        start = 0
        
        for i, v in enumerate(voiced):
            if v and not in_speech:
                in_speech = True
                start = i * hop
            elif not v and in_speech:
                in_speech = False
                end = i * hop
                duration = (end - start) / self.config.sample_rate
                if duration > 0.3:  # 最短300ms
                    segments.append({
                        'start': start,
                        'end': end,
                        'duration_sec': duration,
                    })
        
        return segments
    
    def recognize_emotion(self, audio_segment: np.ndarray) -> dict:
        """
        语音情绪识别
        
        CNN-BiLSTM-Attention 模型
        参考: ETASR 16(3) (2026) "Real-Time Speech Emotion Recognition 
        with CNN-BiLSTM-Attention"
        
        7类情绪：neutral, happy, sad, angry, fear, surprise, disgust
        
        座舱应用：
        - angry → 可能路怒，触发安全提醒
        - sad → 情绪低落，建议休息
        - fear → 紧急情况，主动询问
        - happy → 正常驾驶，保持状态
        """
        # 提取声学特征
        # 1. MFCC
        mfcc = self._extract_mfcc(audio_segment)
        
        # 2. 频谱特征
        spectral_centroid = self._spectral_centroid(audio_segment)
        spectral_rolloff = self._spectral_rolloff(audio_segment)
        
        # 3. 韵律特征
        pitch = self._extract_pitch(audio_segment)
        energy = np.mean(audio_segment**2)
        
        # 情绪分类（实际使用训练好的CNN-BiLSTM-Attention模型）
        emotions = ['neutral', 'happy', 'sad', 'angry', 
                     'fear', 'surprise', 'disgust']
        
        # 简化版规则（实际使用ML模型）
        if energy > 0.5 and np.std(pitch) > 50:
            predicted = 'angry'
            confidence = 0.75
        elif energy < 0.1 and np.mean(pitch) < 120:
            predicted = 'sad'
            confidence = 0.65
        elif np.mean(pitch) > 200 and energy > 0.3:
            predicted = 'happy'
            confidence = 0.70
        else:
            predicted = 'neutral'
            confidence = 0.80
        
        return {
            'emotion': predicted,
            'confidence': confidence,
            'all_scores': {e: np.random.random() for e in emotions},
            'features': {
                'energy': energy,
                'pitch_mean': np.mean(pitch),
                'pitch_std': np.std(pitch),
                'spectral_centroid': spectral_centroid,
            }
        }
    
    def _extract_mfcc(self, audio, n_mfcc=13):
        """提取MFCC特征"""
        # 简化版
        from scipy.fft import dct
        frame = audio[:512]
        spec = np.abs(np.fft.rfft(frame))
        mel_filter = np.random.rand(26, len(spec))  # 概念
        mel_spec = mel_filter @ spec
        mfcc = dct(mel_spec)[:n_mfcc]
        return mfcc
    
    def _spectral_centroid(self, audio):
        spec = np.abs(np.fft.rfft(audio))
        freqs = np.fft.rfftfreq(len(audio), 1/self.config.sample_rate)
        return np.sum(freqs * spec) / (np.sum(spec) + 1e-8)
    
    def _spectral_rolloff(self, audio, roll_percent=0.85):
        spec = np.abs(np.fft.rfft(audio))
        cumsum = np.cumsum(spec)
        rolloff_idx = np.searchsorted(cumsum, roll_percent * cumsum[-1])
        freqs = np.fft.rfftfreq(len(audio), 1/self.config.sample_rate)
        return freqs[rolloff_idx] if rolloff_idx < len(freqs) else freqs[-1]
    
    def _extract_pitch(self, audio):
        """简化版基频提取"""
        autocorr = np.correlate(audio, audio, 'full')
        autocorr = autocorr[len(autocorr)//2:]
        min_lag = int(self.config.sample_rate / 400)  # 最高400Hz
        max_lag = int(self.config.sample_rate / 50)   # 最低50Hz
        
        if max_lag < len(autocorr):
            peak = np.argmax(autocorr[min_lag:max_lag]) + min_lag
            if autocorr[peak] > 0.3 * autocorr[0]:
                return np.array([self.config.sample_rate / peak])
        
        return np.array([0])
    
    def identify_speaker_zone(self, multi_mic_signal: np.ndarray) -> str:
        """
        识别说话人区域
        
        通过到达时间差（TDOA）判断：
        - driver: 方向盘左侧
        - passenger: 前排右侧
        - rear_left: 后排左侧
        - rear_right: 后排右侧
        
        Returns:
            zone: 说话区域
        """
        # 使用GCC-PHAT进行TDOA估计
        mic0 = multi_mic_signal[0]
        mic1 = multi_mic_signal[1] if multi_mic_signal.shape[0] > 1 else mic0
        
        # 互相关
        corr = np.correlate(mic0, mic1, 'full')
        corr = corr[len(corr)//2:]
        
        # 找峰值
        peak_idx = np.argmax(np.abs(corr))
        tdoa_samples = peak_idx
        
        # 时延 → 角度
        c = 343.0
        d = self.config.mic_spacing
        tdoa_sec = tdoa_samples / self.config.sample_rate
        
        if abs(tdoa_sec) < d / c:
            angle = np.degrees(np.arcsin(tdoa_sec * c / d))
        else:
            angle = 0
        
        # 角度 → 区域
        if -60 <= angle <= 60:
            if angle < -15:
                return 'driver'
            elif angle > 15:
                return 'passenger'
            else:
                return 'center_front'
        else:
            return 'rear'
    
    def process_utterance(self, multi_mic_signal: np.ndarray) -> dict:
        """
        完整语音处理管道
        
        Args:
            multi_mic_signal: (num_mics, N) 多麦克风信号
        
        Returns:
            result: 包含转写、情绪、区域等信息
        """
        # 1. 波束形成
        enhanced = self.beamforming(multi_mic_signal, look_direction=0)
        
        # 2. VAD
        segments = self.voice_activity_detection(enhanced)
        
        if not segments:
            return {'status': 'no_speech_detected'}
        
        # 取最长语音段
        longest = max(segments, key=lambda s: s['duration_sec'])
        audio_segment = enhanced[longest['start']:longest['end']]
        
        # 3. 说话人区域识别
        zone = self.identify_speaker_zone(multi_mic_signal)
        
        # 4. 情绪识别
        emotion_result = self.recognize_emotion(audio_segment)
        
        # 5. STT（概念性，实际使用边缘模型）
        transcript = "帮我导航到最近的充电站"  # 概念输出
        
        # 6. 意图理解（概念性）
        intent = "navigation.search"
        entities = {"target": "充电站", "preference": "最近"}
        
        # 7. 响应生成
        if emotion_result['emotion'] == 'angry':
            response = "我理解您可能有些着急。最近的充电站在2公里外，我已经为您规划好路线。"
        elif emotion_result['emotion'] == 'sad':
            response = "好的，为您找到最近的充电站。如果需要休息，前方1公里有服务区。"
        else:
            response = "好的，最近的充电站在前方2公里，已开始导航。"
        
        return {
            'transcript': transcript,
            'intent': intent,
            'entities': entities,
            'emotion': emotion_result['emotion'],
            'emotion_confidence': emotion_result['confidence'],
            'speaker_zone': zone,
            'duration_sec': longest['duration_sec'],
            'response': response,
            'processed_on_device': True,
            'latency_ms': 0,  # 边缘处理
        }


# ==================== 测试 ====================
if __name__ == "__main__":
    np.random.seed(42)
    config = VoiceAgentConfig()
    agent = CabinVoiceAgent(config)
    
    print(f"=== 座舱边缘语音 AI 代理 ===")
    print(f"采样率: {config.sample_rate} Hz")
    print(f"麦克风数: {config.num_mics}")
    print(f"词表大小: {config.vocab_size}")
    print(f"内存限制: {config.max_memory_mb} MB")
    print(f"延迟目标: <{config.max_latency_ms} ms")
    print(f"离线模式: {config.offline_mode}")
    
    # 模拟4麦克风信号（3秒）
    duration = 3.0
    N = int(config.sample_rate * duration)
    t = np.arange(N) / config.sample_rate
    
    # 模拟语音信号（含噪声）
    speech = 0.3 * np.sin(2*np.pi*200*t) * np.exp(-((t-1.5)**2)/0.5)
    engine_noise = 0.1 * np.random.normal(0, 1, N)
    
    # 4麦克风（不同延迟）
    multi_mic = np.zeros((config.num_mics, N))
    for i in range(config.num_mics):
        delay = i * 2  # 2 samples delay
        multi_mic[i] = np.roll(speech + engine_noise, delay)
    
    # 处理
    result = agent.process_utterance(multi_mic)
    
    print(f"\n=== 处理结果 ===")
    print(f"转写: {result.get('transcript', 'N/A')}")
    print(f"意图: {result.get('intent', 'N/A')}")
    print(f"情绪: {result.get('emotion', 'N/A')}")
    print(f"情绪置信度: {result.get('emotion_confidence', 0):.2f}")
    print(f"说话区域: {result.get('speaker_zone', 'N/A')}")
    print(f"语音时长: {result.get('duration_sec', 0):.1f}秒")
    print(f"响应: {result.get('response', 'N/A')}")
    print(f"边缘处理: {result.get('processed_on_device', False)}")
    
    # 成本对比
    print(f"\n=== 成本对比 ===")
    # STT API: $0.006/min, 2min/day, 120K vehicles
    stt_cost_per_vehicle = 0.006 * 2 * 365  # $4.38/year
    fleet_size = 120000
    
    print(f"云端STT成本: ${stt_cost_per_vehicle:.2f}/年/车")
    print(f"10万车队/年: ${stt_cost_per_vehicle * fleet_size:.0f}")
    print(f"10年累积: ${stt_cost_per_vehicle * fleet_size * 10:.0f}")
    print(f"边缘STT成本: $0（一次性硬件）")
    print(f"节省(10年/10万车): ${stt_cost_per_vehicle * fleet_size * 10:.0f}")
    
    print(f"\n=== 竞品对比 ===")
    competitors = [
        ('Mercedes-Benz × Liquid AI', 'MB.OS', '2026 H2', '全栈NLP'),
        ('Hyundai Gleo AI', 'Grandeur', '已量产', '上下文感知'),
        ('Lucid × SoundHound', 'Lucid', '已部署', '离线对话'),
        ('BMW × Alexa', 'Panoramic iDrive', '已量产', '混合云'),
        ('Tesla Voice', 'FSD', '已量产', '云端优先'),
    ]
    
    for name, platform, status, feature in competitors:
        print(f"  {name}: {platform} ({status}, {feature})")

"""
CNN-BiLSTM-Attention 语音情绪识别
参考: ETASR 16(3) (2026)

架构：
1. CNN: 从MFCC提取局部声学特征
2. BiLSTM: 捕捉时序依赖
3. Attention: 聚焦情绪关键帧
4. FC: 7类情绪分类

精度指标：
- IEMOCAP 4类：76.3% (val-test)
- 车内噪声环境：68.5% (-8%)
- 实时推理：<50ms（量化模型）
"""

import torch
import torch.nn as nn

class CNNBiLSTMAttention(nn.Module):
    """
    语音情绪识别模型
    
    输入: (B, 1, T, 39) MFCC + delta + delta-delta
    输出: (B, 7) 情绪类别概率
    """
    
    def __init__(self, num_emotions=7, input_dim=39):
        super().__init__()
        
        # CNN 特征提取
        self.cnn = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d((2, 1)),
            
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d((2, 1)),
            
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d((2, 1)),
        )
        
        # BiLSTM 时序建模
        self.lstm = nn.LSTM(
            input_size=128 * (input_dim // 8),
            hidden_size=128,
            num_layers=2,
            batch_first=True,
            bidirectional=True
        )
        
        # Attention
        self.attention = nn.Sequential(
            nn.Linear(256, 64),
            nn.Tanh(),
            nn.Linear(64, 1),
            nn.Softmax(dim=1)
        )
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64, num_emotions)
        )
    
    def forward(self, x):
        # x: (B, 1, T, 39)
        B = x.size(0)
        
        # CNN
        cnn_out = self.cnn(x)  # (B, 128, T//8, 39//8)
        cnn_out = cnn_out.view(B, cnn_out.size(2), -1)  # (B, T//8, 128*4)
        
        # BiLSTM
        lstm_out, _ = self.lstm(cnn_out)  # (B, T//8, 256)
        
        # Attention
        att_weights = self.attention(lstm_out)  # (B, T//8, 1)
        context = torch.sum(att_weights * lstm_out, dim=1)  # (B, 256)
        
        # 分类
        output = self.classifier(context)  # (B, 7)
        
        return output, att_weights.squeeze(-1)