import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
import numpy as np

class TextGuidedSER(nn.Module):
    """
    文本引导环境感知语音情绪识别
    
    论文: Leem et al., arXiv 2407.17716
    核心: 用环境文本描述引导SER模型适应噪声
    
    车载场景应用:
    - 环境描述: "汽车座舱内，引擎噪声和风噪"
    - SNR范围: 0-15dB（典型座舱）
    - 情绪维度: arousal(唤醒) / valence(效价) / dominance(支配)
    """
    
    def __init__(self, speech_encoder_name='facebook/wav2vec2-base',
                 text_encoder_name='bert-base-uncased',
                 emotion_dim=3,  # arousal, valence, dominance
                 hidden_dim=256):
        super().__init__()
        
        # 语音编码器（Transformer-based SSL）
        self.speech_encoder = AutoModel.from_pretrained(speech_encoder_name)
        speech_feat_dim = self.speech_encoder.config.hidden_size  # 768
        
        # 文本编码器（LLM/CLAP）
        self.text_tokenizer = AutoTokenizer.from_pretrained(text_encoder_name)
        self.text_encoder = AutoModel.from_pretrained(text_encoder_name)
        text_feat_dim = self.text_encoder.config.hidden_size  # 768
        
        # 融合层
        self.fusion = nn.Sequential(
            nn.Linear(speech_feat_dim + text_feat_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.LayerNorm(hidden_dim // 2),
            nn.GELU(),
            nn.Dropout(0.1),
        )
        
        # 情绪回归头
        self.emotion_head = nn.Linear(hidden_dim // 2, emotion_dim)
        
    def forward(self, speech, text_input_ids, text_attention_mask):
        """
        前向传播
        
        Args:
            speech: 语音波形, shape=(B, T)
            text_input_ids: 环境描述token, shape=(B, L)
            text_attention_mask: 文本attention mask
            
        Returns:
            emotions: (arousal, valence, dominance), shape=(B, 3)
        """
        # 语音特征
        speech_outputs = self.speech_encoder(speech)
        speech_features = speech_outputs.last_hidden_state.mean(dim=1)  # (B, 768)
        
        # 文本特征（环境描述）
        text_outputs = self.text_encoder(
            input_ids=text_input_ids,
            attention_mask=text_attention_mask
        )
        text_features = text_outputs.last_hidden_state[:, 0, :]  # CLS token, (B, 768)
        
        # 融合
        combined = torch.cat([speech_features, text_features], dim=-1)
        fused = self.fusion(combined)
        
        # 情绪预测
        emotions = self.emotion_head(fused)
        
        return emotions
    
    def predict_emotion(self, speech, environment_text):
        """
        预测情绪（推理时使用）
        
        Args:
            speech: 语音波形
            environment_text: 环境描述，如"在汽车座舱内，引擎噪声"
            
        Returns:
            dict: {arousal, valence, dominance}
        """
        # 编码环境文本
        tokens = self.text_tokenizer(
            environment_text,
            return_tensors='pt',
            padding=True,
            truncation=True,
            max_length=128
        )
        
        with torch.no_grad():
            emotions = self.forward(
                speech.unsqueeze(0),
                tokens['input_ids'],
                tokens['attention_mask']
            )
        
        return {
            'arousal': emotions[0, 0].item(),    # 1-5
            'valence': emotions[0, 1].item(),     # 1-5
            'dominance': emotions[0, 2].item()    # 1-5
        }


# 训练代码
class TG_EAT_Trainer:
    """
    文本引导环境感知训练器
    
    训练策略:
    1. 干净语音 + 环境噪声 → 带噪语音
    2. 环境文本描述 → 文本嵌入
    3. 带噪语音 + 文本嵌入 → 情绪预测
    4. 损失: MSE (arousal/valence/dominance)
    """
    
    # 车载常见噪声环境模板
    CABIN_NOISE_TEMPLATES = [
        "Inside a car cabin with engine noise and road noise",
        "Inside a vehicle driving on highway with wind noise",
        "Inside a car with air conditioning running",
        "Inside an electric vehicle with minimal motor noise",
        "Inside a car with rain noise on the roof",
        "Inside a truck cabin with diesel engine noise",
        "Inside a car with music playing in background",
        "Inside a vehicle in heavy traffic with horn sounds",
    ]
    
    def __init__(self, model, learning_rate=1e-5):
        self.model = model
        self.optimizer = torch.optim.AdamW(
            model.parameters(), 
            lr=learning_rate,
            weight_decay=0.01
        )
        self.criterion = nn.MSELoss()
        
    def add_noise(self, clean_speech, noise, snr_db):
        """
        添加噪声到干净语音
        
        Args:
            clean_speech: 干净语音
            noise: 噪声信号
            snr_db: 目标信噪比(dB)
        """
        speech_power = torch.mean(clean_speech ** 2)
        noise_power = torch.mean(noise ** 2)
        
        # 计算缩放因子
        snr_linear = 10 ** (snr_db / 10)
        scale = torch.sqrt(speech_power / (noise_power * snr_linear + 1e-8))
        
        noisy_speech = clean_speech + scale * noise
        return noisy_speech
    
    def train_step(self, clean_speech, noise, emotion_labels, 
                   env_text_ids, env_text_mask, snr_db=None):
        """
        单步训练
        
        Args:
            clean_speech: 干净语音 (B, T)
            noise: 噪声信号 (B, T)
            emotion_labels: 情绪标签 (B, 3) [arousal, valence, dominance]
            env_text_ids: 环境文本token
            env_text_mask: 文本mask
            snr_db: 信噪比，随机采样
        """
        if snr_db is None:
            snr_db = np.random.choice([-5, 0, 5, 10, 15, 20])
        
        # 添加噪声
        noisy_speech = self.add_noise(clean_speech, noise, snr_db)
        
        # 前向传播
        pred_emotions = self.model(noisy_speech, env_text_ids, env_text_mask)
        
        # 损失
        loss = self.criterion(pred_emotions, emotion_labels)
        
        # 反向传播
        self.optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
        self.optimizer.step()
        
        return loss.item(), snr_db


# 车载场景应用示例
if __name__ == "__main__":
    print("=== 车载语音情绪识别噪声鲁棒性测试 ===\n")
    
    # 模型初始化
    model = TextGuidedSER()
    model.eval()
    
    # 模拟带噪语音
    np.random.seed(42)
    speech = torch.randn(16000 * 3)  # 3秒语音
    
    # 不同环境描述
    environments = [
        "Inside a quiet electric vehicle cabin",
        "Inside a car on highway with heavy wind noise", 
        "Inside a truck with diesel engine noise",
        "Inside a sports car with loud engine noise",
    ]
    
    print("环境描述对情绪识别的影响：\n")
    for env in environments:
        result = model.predict_emotion(speech, env)
        print(f"环境: {env}")
        print(f"  Arousal: {result['arousal']:.2f} (唤醒度)")
        print(f"  Valence: {result['valence']:.2f} (效价)")
        print(f"  Dominance: {result['dominance']:.2f} (支配度)")
        print()
    
    # 论文性能数据
    print("=== 论文报告性能（-5dB SNR）===")
    print("vs 最佳基线:")
    print("  Arousal:   +31.8%")
    print("  Dominance: +23.5%")
    print("  Valence:   +9.5%")
    print()
    print("vs 原始SSL模型:")
    print("  Arousal:   +163.6%")
    print("  Dominance: +200.0%")
    print("  Valence:   +91.6%")

class CabinEnvironmentGenerator:
    """
    自动生成座舱环境文本描述
    
    输入: GPS速度、车速、空调状态、车窗状态、雨刮状态
    输出: 环境文本描述（用于TG-EAT）
    """
    
    TEMPLATES = {
        'ev_cruise': "Inside an electric vehicle cruising at {speed} km/h with wind noise",
        'city_traffic': "Inside a car in city traffic with engine noise and intermittent horn sounds",
        'highway': "Inside a car on highway at {speed} km/h with steady wind and tire noise",
        'rain': "Inside a car with rain noise on the roof and wipers running",
        'ac_on': "Inside a car with air conditioning fan noise",
        'window_open': "Inside a car with window open, strong wind noise",
        'tunnel': "Inside a tunnel with echo and low-frequency resonance",
        'idle': "Inside a parked car with engine idling",
    }
    
    def generate(self, speed=0, ev=False, ac=False, window=False, 
                 rain=False, tunnel=False):
        """生成环境描述"""
        if tunnel:
            return self.TEMPLATES['tunnel']
        if rain:
            return self.TEMPLATES['rain'].format(speed=speed)
        if window:
            return self.TEMPLATES['window_open']
        if speed > 80:
            template = self.TEMPLATES['highway'].format(speed=speed)
        elif speed > 20:
            template = self.TEMPLATES['city_traffic']
        elif ev and speed > 0:
            template = self.TEMPLATES['ev_cruise'].format(speed=speed)
        else:
            template = self.TEMPLATES['idle']
        
        if ac:
            template += " and air conditioning running"
        
        return template