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

class ModalityEncoder(nn.Module):
    """单模态编码器基类"""
    
    def __init__(self, input_dim: int, hidden_dim: int = 128):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU()
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.encoder(x)


class VisualEncoder(ModalityEncoder):
    """
    视觉模态编码器
    
    输入特征向量：
    - gaze_xy: (2,) 视线方向
    - gaze_confidence: (1,) 置信度
    - head_pose: (6,) 头部姿态 (pitch, yaw, roll + 角速度)
    - body_activity: (4,) 身体活动量 (4个身体部位)
    - interaction_score: (3,) 物体交互评分 (手机/中控/方向盘)
    
    总输入维度: 16
    """
    
    def __init__(self, hidden_dim: int = 128):
        super().__init__(input_dim=16, hidden_dim=hidden_dim)
        
        # 时序建模
        self.temporal = nn.LSTM(
            input_size=hidden_dim,
            hidden_size=hidden_dim,
            num_layers=2,
            batch_first=True,
            bidirectional=False
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, T, 16) 时序视觉特征
        Returns:
            encoded: (B, hidden_dim) 模态级嵌入
        """
        B, T, _ = x.shape
        feat = self.encoder(x.reshape(B * T, -1))
        feat = feat.reshape(B, T, -1)
        _, (hn, _) = self.temporal(feat)
        return hn[-1]  # 最后一层最后时刻步


class AudioEncoder(ModalityEncoder):
    """
    音频模态编码器
    
    输入特征向量：
    - scene_embedding: (8,) 声学场景嵌入
    - speech_features: (6,) 语音韵律 (语速/停顿/音高/能量)
    - breathing_rate: (2,) 呼吸频率和方差
    - attention_score: (1,) 注意力声学评分
    
    总输入维度: 17
    """
    
    def __init__(self, hidden_dim: int = 128):
        super().__init__(input_dim=17, hidden_dim=hidden_dim)
        self.temporal = nn.LSTM(
            input_size=hidden_dim, hidden_size=hidden_dim,
            num_layers=1, batch_first=True
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, _ = x.shape
        feat = self.encoder(x.reshape(B * T, -1)).reshape(B, T, -1)
        _, (hn, _) = self.temporal(feat)
        return hn[-1]


class VitalSignEncoder(ModalityEncoder):
    """
    生命体征模态编码器
    
    输入特征向量：
    - heart_rate: (2,) 心率和变化率
    - hrv_features: (5,) HRV时域频域特征
    - eda_features: (3,) 皮肤电活动特征
    - respiration: (2,) 呼吸率和方差
    
    总输入维度: 12
    """
    
    def __init__(self, hidden_dim: int = 128):
        super().__init__(input_dim=12, hidden_dim=hidden_dim)
        self.temporal = nn.LSTM(
            input_size=hidden_dim, hidden_size=hidden_dim,
            num_layers=1, batch_first=True
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, _ = x.shape
        feat = self.encoder(x.reshape(B * T, -1)).reshape(B, T, -1)
        _, (hn, _) = self.temporal(feat)
        return hn[-1]


class VehicleDataEncoder(ModalityEncoder):
    """
    车辆数据模态编码器
    
    输入特征向量：
    - steering_micro: (3,) 方向盘微操作 (角度/速度/加速度)
    - pedal_pattern: (2,) 踏板操作平稳度
    - lane_keeping: (2,) 车道偏移和方差
    - headway: (2,) 前车距离和变化率
    
    总输入维度: 9
    """
    
    def __init__(self, hidden_dim: int = 128):
        super().__init__(input_dim=9, hidden_dim=hidden_dim)
        self.temporal = nn.LSTM(
            input_size=hidden_dim, hidden_size=hidden_dim,
            num_layers=1, batch_first=True
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, _ = x.shape
        feat = self.encoder(x.reshape(B * T, -1)).reshape(B, T, -1)
        _, (hn, _) = self.temporal(feat)
        return hn[-1]


class SensAIFusionModel(nn.Module):
    """
    SensAI 多模态融合模型
    
    架构：
    1. 四个专用编码器提取各模态嵌入
    2. Cross-Attention 融合层实现跨模态交互
    3. 认知状态解码器输出认知分心评估
    
    输出：
    - cognitive_state: (B, 3) [正常, 轻度分心, 严重分心] 概率
    - contributing_factors: (B, 4) 各模态贡献度
    """
    
    def __init__(self, hidden_dim: int = 128, num_heads: int = 4):
        super().__init__()
        
        # 模态编码器
        self.visual_enc = VisualEncoder(hidden_dim)
        self.audio_enc = AudioEncoder(hidden_dim)
        self.vital_enc = VitalSignEncoder(hidden_dim)
        self.vehicle_enc = VehicleDataEncoder(hidden_dim)
        
        # Cross-Attention 融合
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=hidden_dim,
            num_heads=num_heads,
            batch_first=True
        )
        
        # 模态权重学习
        self.modality_weights = nn.Sequential(
            nn.Linear(hidden_dim * 4, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 4),
            nn.Softmax(dim=-1)
        )
        
        # 认知状态解码器
        self.state_decoder = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, 3),  # 3类认知状态
            nn.Softmax(dim=-1)
        )
        
        # 贡献因素解释器
        self.factor_explainer = nn.Sequential(
            nn.Linear(hidden_dim * 4, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 4),
            nn.Sigmoid()
        )
    
    def forward(self, visual: torch.Tensor, audio: torch.Tensor,
                vital: torch.Tensor, vehicle: torch.Tensor) -> dict:
        """
        Args:
            visual: (B, T, 16)
            audio: (B, T, 17)
            vital: (B, T, 12)
            vehicle: (B, T, 9)
        
        Returns:
            dict with 'cognitive_state', 'contributing_factors'
        """
        # 各模态编码
        v_emb = self.visual_enc(visual)    # (B, hidden_dim)
        a_emb = self.audio_enc(audio)      # (B, hidden_dim)
        s_emb = self.vital_enc(vital)     # (B, hidden_dim)
        c_emb = self.vehicle_enc(vehicle) # (B, hidden_dim)
        
        # 堆叠为序列
        stacked = torch.stack([v_emb, a_emb, s_emb, c_emb], dim=1)  # (B, 4, hidden)
        
        # Cross-Attention
        fused, _ = self.cross_attn(stacked, stacked, stacked)
        fused_repr = fused.mean(dim=1)  # (B, hidden_dim)
        
        # 模态权重
        concat_emb = torch.cat([v_emb, a_emb, s_emb, c_emb], dim=-1)
        weights = self.modality_weights(concat_emb)  # (B, 4)
        
        # 加权融合
        weighted_emb = (stacked * weights.unsqueeze(-1)).sum(dim=1)
        final_repr = fused_repr + weighted_emb
        
        # 解码
        cognitive_state = self.state_decoder(final_repr)
        contributing_factors = self.factor_explainer(concat_emb)
        
        return {
            'cognitive_state': cognitive_state,
            'contributing_factors': contributing_factors,
            'modality_weights': weights
        }


# 测试
if __name__ == "__main__":
    model = SensAIFusionModel(hidden_dim=64, num_heads=4)
    
    B, T = 2, 30  # batch=2, 30帧序列
    
    # 模拟输入
    visual = torch.randn(B, T, 16)
    audio = torch.randn(B, T, 17)
    vital = torch.randn(B, T, 12)
    vehicle = torch.randn(B, T, 9)
    
    output = model(visual, audio, vital, vehicle)
    
    print("=== SensAI 融合模型输出 ===")
    print(f"认知状态: {output['cognitive_state'].shape}")
    print(f"  正常概率: {output['cognitive_state'][0, 0]:.3f}")
    print(f"  轻度分心: {output['cognitive_state'][0, 1]:.3f}")
    print(f"  严重分心: {output['cognitive_state'][0, 2]:.3f}")
    print(f"\n模态权重: {output['modality_weights'][0].tolist()}")
    print(f"  视觉: {output['modality_weights'][0, 0]:.3f}")
    print(f"  音频: {output['modality_weights'][0, 1]:.3f}")
    print(f"  生命体征: {output['modality_weights'][0, 2]:.3f}")
    print(f"  车辆: {output['modality_weights'][0, 3]:.3f}")
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"\n模型参数量: {total_params:,}")
    print(f"模型大小: {total_params * 4 / 1024 / 1024:.1f} MB (FP32)")

class ContextSensitiveWarning:
    """
    SensAI 上下文敏感警告系统
    
    与传统DMS的区别：
    - 传统：PERCLOS > 30% → 警告（不看场景）
    - SensAI：PERCLOS > 30% + 高速公路 + 直道 → 不警告
              PERCLOS > 30% + 城市路口 + 行人 → 紧急警告
    
    避免不必要或不当的警告
    """
    
    SCENE_WEIGHTS = {
        'highway_straight': 0.3,      # 低风险场景
        'highway_traffic': 0.6,        # 中风险
        'urban_intersection': 1.0,     # 高风险
        'urban_pedestrian': 1.2,       # 最高风险
        'parking_lot': 0.2,           # 低风险
        'night_rural': 0.5,           # 中风险
    }
    
    def evaluate(self, cognitive_state: dict, 
                 driving_scene: str,
                 vehicle_speed: float) -> dict:
        """
        评估是否需要发出警告
        
        Args:
            cognitive_state: {'normal': float, 'mild': float, 'severe': float}
            driving_scene: 场景类型
            vehicle_speed: km/h
            
        Returns:
            {'should_warn': bool, 'level': str, 'reason': str}
        """
        scene_weight = self.SCENE_WEIGHTS.get(driving_scene, 0.5)
        
        # 速度调整因子
        speed_factor = min(vehicle_speed / 60.0, 1.5)
        
        # 综合风险评分
        risk_score = (
            cognitive_state.get('mild', 0) * 0.3 +
            cognitive_state.get('severe', 0) * 1.0
        ) * scene_weight * speed_factor
        
        if risk_score > 0.7:
            return {'should_warn': True, 'level': 'urgent', 
                    'reason': f'高风险场景({driving_scene})+认知分心',
                    'risk_score': risk_score}
        elif risk_score > 0.4:
            return {'should_warn': True, 'level': 'moderate',
                    'reason': f'中风险：{driving_scene}场景下的认知分心',
                    'risk_score': risk_score}
        elif risk_score > 0.2:
            return {'should_warn': True, 'level': 'mild',
                    'reason': '轻度认知分心提示',
                    'risk_score': risk_score}
        else:
            return {'should_warn': False, 'level': 'none',
                    'reason': '认知状态正常或场景风险低',
                    'risk_score': risk_score}


# 测试
if __name__ == "__main__":
    warning_sys = ContextSensitiveWarning()
    
    test_cases = [
        # (认知状态, 驾驶场景, 车速)
        ({'normal': 0.7, 'mild': 0.2, 'severe': 0.1}, 'highway_straight', 90),
        ({'normal': 0.3, 'mild': 0.4, 'severe': 0.3}, 'urban_intersection', 30),
        ({'normal': 0.2, 'mild': 0.3, 'severe': 0.5}, 'urban_pedestrian', 20),
        ({'normal': 0.8, 'mild': 0.15, 'severe': 0.05}, 'parking_lot', 10),
    ]
    
    for state, scene, speed in test_cases:
        result = warning_sys.evaluate(state, scene, speed)
        print(f"场景={scene}, 速度={speed}km/h → {result['level']} ({result['risk_score']:.2f})")

class PracticalCognitiveDistractionDetector:
    """
    面向量产的认知分心检测方案（Phase 1-2）
    
    仅使用现有DMS摄像头 + CAN总线数据
    无需额外传感器
    
    核心特征：
    1. 眼动熵（gaze entropy）— 认知分心时降低
    2. 扫视频率（saccade rate）— 认知分心时降低
    3. 方向盘微操作频率 — 认知分心时降低
    4. 车道保持标准差 — 认知分心时增加
    5. 反应延迟代理 — 眼动变化率
    """
    
    def __init__(self, fps: int = 30):
        self.fps = fps
        self.window = 30 * fps  # 30秒窗口
        self.baselines = {
            'gaze_entropy': 2.5,
            'saccade_rate': 3.0,
            'steer_micro_rate': 0.5,  # Hz
            'lane_dev_std': 0.15,     # m
        }
    
    def detect(self, gaze_sequence: np.ndarray,
               steering_data: np.ndarray,
               lane_data: np.ndarray) -> dict:
        """
        从现有DMS和CAN数据检测认知分心
        
        Args:
            gaze_sequence: (N, 2) 视线序列
            steering_data: (N,) 方向盘角度序列
            lane_data: (N,) 车道偏移序列
            
        Returns:
            {'cognitive_distraction': float, 'confidence': float, 'indicators': dict}
        """
        indicators = {}
        
        # 1. 眼动熵
        hist, _, _ = np.histogram2d(
            gaze_sequence[:, 0], gaze_sequence[:, 1], bins=5
        )
        hist_norm = hist / (hist.sum() + 1e-8)
        indicators['gaze_entropy'] = -np.sum(
            hist_norm * np.log(hist_norm + 1e-8)
        )
        
        # 2. 扫视频率
        gaze_diff = np.diff(gaze_sequence, axis=0)
        gaze_speed = np.linalg.norm(gaze_diff, axis=1)
        indicators['saccade_rate'] = np.sum(
            gaze_speed > 0.15
        ) / (len(gaze_sequence) / self.fps)
        
        # 3. 方向盘微操作
        steer_diff = np.diff(steering_data)
        micro_corrections = np.sum(np.abs(steer_diff) < 0.5) 
        indicators['steer_micro_rate'] = micro_corrections / (
            len(steering_data) / self.fps
        )
        
        # 4. 车道保持
        indicators['lane_dev_std'] = np.std(lane_data)
        
        # 综合评分
        deviations = {}
        for key, val in indicators.items():
            baseline = self.baselines.get(key, 1.0)
            deviations[key] = abs(val - baseline) / baseline
        
        score = np.mean(list(deviations.values()))
        
        return {
            'cognitive_distraction': float(min(score, 1.0)),
            'confidence': float(min(score * 1.5, 0.95)),
            'indicators': indicators,
            'deviations': deviations
        }


# 测试
if __name__ == "__main__":
    detector = PracticalCognitiveDistractionDetector(fps=30)
    
    N = 900  # 30秒数据
    
    # 模拟正常驾驶
    normal_gaze = np.column_stack([
        np.random.uniform(-0.4, 0.4, N),
        np.random.uniform(-0.2, 0.2, N)
    ])
    normal_steer = np.cumsum(np.random.normal(0, 0.3, N))
    normal_lane = np.random.normal(0, 0.1, N)
    
    # 模拟认知分心（眼动范围收窄、方向盘微操作减少）
    distracted_gaze = np.column_stack([
        np.random.uniform(-0.1, 0.1, N),
        np.random.uniform(-0.05, 0.05, N)
    ])
    distracted_steer = np.cumsum(np.random.normal(0, 0.5, N))  # 大幅修正多，微修正少
    distracted_lane = np.random.normal(0, 0.3, N)  # 车道偏移增大
    
    normal_result = detector.detect(normal_gaze, normal_steer, normal_lane)
    distracted_result = detector.detect(distracted_gaze, distracted_steer, distracted_lane)
    
    print("=== 正常驾驶 ===")
    print(f"认知分心评分: {normal_result['cognitive_distraction']:.3f}")
    print(f"眼动熵: {normal_result['indicators']['gaze_entropy']:.3f}")
    print(f"扫视频率: {normal_result['indicators']['saccade_rate']:.2f}/s")
    
    print("\n=== 认知分心 ===")
    print(f"认知分心评分: {distracted_result['cognitive_distraction']:.3f}")
    print(f"眼动熵: {distracted_result['indicators']['gaze_entropy']:.3f}")
    print(f"扫视频率: {distracted_result['indicators']['saccade_rate']:.2f}/s")