import torch
import torch.nn as nn
from typing import Dict, List

class SensAIEncoder(nn.Module):
    """
    SensAI多模态编码器基类
    
    设计原则：
    1. 每个模态有独立编码器
    2. 编码器输出标准化表示（统一维度）
    3. 支持时序建模（LSTM/Transformer）
    4. 缺失模态时仍可工作（降级运行）
    """
    
    def __init__(self, input_dim: int, hidden_dim: int = 256,
                 num_layers: int = 2, dropout: float = 0.1):
        super().__init__()
        self.hidden_dim = hidden_dim
        
        # 时序编码器
        self.lstm = nn.LSTM(
            input_size=input_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout
        )
        
        # 注意力池化
        self.attention = nn.MultiheadAttention(
            embed_dim=hidden_dim,
            num_heads=4,
            batch_first=True
        )
        
        # 模态可靠性估计
        self.reliability_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.GELU(),
            nn.Linear(hidden_dim // 2, 1),
            nn.Sigmoid()
        )
    
    def forward(self, x: torch.Tensor, 
                mask: torch.Tensor = None) -> Dict[str, torch.Tensor]:
        """
        Args:
            x: (B, T, input_dim) 模态输入序列
            mask: (B, T) 有效帧掩码
        
        Returns:
            encoded: (B, hidden_dim) 编码特征
            reliability: (B, 1) 模态可靠性
        """
        # LSTM编码
        lstm_out, _ = self.lstm(x)  # (B, T, hidden)
        
        # 自注意力池化
        attn_out, _ = self.attention(lstm_out, lstm_out, lstm_out,
                                      key_padding_mask=mask)
        
        # 时序池化（取最后帧的注意力输出）
        encoded = attn_out.mean(dim=1)  # (B, hidden)
        
        # 模态可靠性
        reliability = self.reliability_head(encoded)  # (B, 1)
        
        return {
            'encoded': encoded,
            'reliability': reliability,
            'temporal_features': lstm_out
        }


class VisualEncoder(SensAIEncoder):
    """视觉编码器：视线+姿态+活动"""
    
    def __init__(self, hidden_dim: int = 256):
        super().__init__(input_dim=128, hidden_dim=hidden_dim)
        
        # 视线特征投影
        self.gaze_proj = nn.Linear(4, 32)  # (pitch, yaw, confidence, velocity)
        
        # 姿态特征投影
        self.pose_proj = nn.Linear(34, 64)  # 17 keypoints × 2D
        
        # 活动特征投影
        self.activity_proj = nn.Linear(10, 32)  # 活动one-hot + confidence
        

class AcousticEncoder(SensAIEncoder):
    """声学编码器：场景分析+语音情感"""
    
    def __init__(self, hidden_dim: int = 256):
        super().__init__(input_dim=64, hidden_dim=hidden_dim)
        
        # 声学特征投影
        self.audio_proj = nn.Linear(64, 64)  # MFCC + 情感特征
        

class PhysiologicalEncoder(SensAIEncoder):
    """生理编码器：HRV+呼吸频率"""
    
    def __init__(self, hidden_dim: int = 256):
        super().__init__(input_dim=8, hidden_dim=hidden_dim)
        
        # rPPG特征投影
        self.rppg_proj = nn.Linear(8, 32)  # HR, HRV, 呼吸率等
        

class VehicleEncoder(SensAIEncoder):
    """车辆数据编码器：方向盘+踏板+加速度"""
    
    def __init__(self, hidden_dim: int = 256):
        super().__init__(input_dim=12, hidden_dim=hidden_dim)
        
        # CAN信号投影
        self.can_proj = nn.Linear(12, 32)  # 方向盘角度+速度+加速度等


class SensAIFusionModel(nn.Module):
    """
    SensAI多模态融合模型
    
    将4个模态编码器的输出融合为认知状态数字孪生
    
    架构：
    1. 4个独立编码器
    2. 可靠性加权融合
    3. 认知状态解码
    4. 上下文敏感警告触发
    """
    
    def __init__(self, hidden_dim: int = 256, 
                 num_distraction_levels: int = 4):
        super().__init__()
        
        # 4个模态编码器
        self.visual_encoder = VisualEncoder(hidden_dim)
        self.acoustic_encoder = AcousticEncoder(hidden_dim)
        self.physio_encoder = PhysiologicalEncoder(hidden_dim)
        self.vehicle_encoder = VehicleEncoder(hidden_dim)
        
        # 模态融合层
        self.fusion = nn.Sequential(
            nn.Linear(hidden_dim * 4, hidden_dim * 2),
            nn.GELU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.GELU(),
            nn.Dropout(0.2)
        )
        
        # 认知状态解码
        self.state_decoder = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, num_distraction_levels)
        )
        
        # 影响因子分析
        self.factor_analyzer = nn.Linear(hidden_dim, 4)  # 4个模态的贡献度
    
    def forward(self, visual, acoustic, physio, vehicle,
                masks: Dict = None) -> Dict[str, torch.Tensor]:
        """
        Args:
            visual: (B, T, 128) 视觉特征序列
            acoustic: (B, T, 64) 声学特征序列
            physio: (B, T, 8) 生理特征序列
            vehicle: (B, T, 12) 车辆特征序列
            masks: 各模态的有效帧掩码
        
        Returns:
            distraction_logits: (B, 4) 分心等级
            distraction_level: (B,) 预测等级
            modality_weights: (B, 4) 各模态贡献度
        """
        # 编码各模态
        vis_out = self.visual_encoder(visual, 
                                      masks.get('visual') if masks else None)
        aco_out = self.acoustic_encoder(acoustic,
                                        masks.get('acoustic') if masks else None)
        phy_out = self.physio_encoder(physio,
                                      masks.get('physio') if masks else None)
        veh_out = self.vehicle_encoder(vehicle,
                                       masks.get('vehicle') if masks else None)
        
        # 可靠性加权
        vis_weighted = vis_out['encoded'] * vis_out['reliability']
        aco_weighted = aco_out['encoded'] * aco_out['reliability']
        phy_weighted = phy_out['encoded'] * phy_out['reliability']
        veh_weighted = veh_out['encoded'] * veh_out['reliability']
        
        # 拼接
        concat = torch.cat([
            vis_weighted, aco_weighted, 
            phy_weighted, veh_weighted
        ], dim=-1)
        
        # 融合
        fused = self.fusion(concat)
        
        # 解码
        logits = self.state_decoder(fused)
        level = logits.argmax(dim=-1)
        
        # 影响因子
        weights = self.factor_analyzer(fused)
        weights = torch.softmax(weights, dim=-1)
        
        return {
            'distraction_logits': logits,
            'distraction_level': level,
            'modality_weights': weights,
            'reliabilities': {
                'visual': vis_out['reliability'],
                'acoustic': aco_out['reliability'],
                'physio': phy_out['reliability'],
                'vehicle': veh_out['reliability']
            }
        }


# 测试
if __name__ == "__main__":
    print("=" * 60)
    print("SensAI 多模态认知分心检测 - 模拟测试")
    print("=" * 60)
    
    model = SensAIFusionModel()
    total_params = sum(p.numel() for p in model.parameters())
    print(f"\n模型参数量: {total_params:,}")
    
    # 模拟输入（30帧 = 1秒@30fps）
    B, T = 2, 30
    visual = torch.randn(B, T, 128)
    acoustic = torch.randn(B, T, 64)
    physio = torch.randn(B, T, 8)
    vehicle = torch.randn(B, T, 12)
    
    # 前向传播
    with torch.no_grad():
        output = model(visual, acoustic, physio, vehicle)
    
    print(f"\n输入:")
    print(f"  视觉: {visual.shape}")
    print(f"  声学: {acoustic.shape}")
    print(f"  生理: {physio.shape}")
    print(f"  车辆: {vehicle.shape}")
    
    print(f"\n输出:")
    print(f"  分心等级logits: {output['distraction_logits'].shape}")
    print(f"  预测等级: {output['distraction_level']}")
    print(f"  模态权重: {output['modality_weights']}")
    
    # 模态降级测试
    print(f"\n{'='*60}")
    print("模态降级测试（生理信号缺失）")
    print(f"{'='*60}")
    
    physio_mask = torch.ones(B, T)
    physio_mask[:, :] = 0  # 标记生理信号全部无效
    
    with torch.no_grad():
        output_degraded = model(visual, acoustic, physio, vehicle,
                                masks={'physio': physio_mask.bool()})
    
    print(f"  视觉可靠性: {output_degraded['reliabilities']['visual'].mean():.3f}")
    print(f"  声学可靠性: {output_degraded['reliabilities']['acoustic'].mean():.3f}")
    print(f"  生理可靠性: {output_degraded['reliabilities']['physio'].mean():.3f}")
    print(f"  车辆可靠性: {output_degraded['reliabilities']['vehicle'].mean():.3f}")
    
    print(f"\n  降级时模态权重变化:")
    print(f"    正常: {output['modality_weights'][0]}")
    print(f"    降级: {output_degraded['modality_weights'][0]}")