import numpy as np
import torch
import torch.nn as nn
from typing import Dict, Tuple

class MultiModalAttentionDetector(nn.Module):
    """
    多模态认知分心检测模型
    
    融合EEG + 眼动 + 驾驶行为三模态
    
    论文核心思路：
    1. 每个模态单独编码
    2. 跨模态注意力融合
    3. 时序建模检测认知状态变化
    """
    
    def __init__(self, 
                 eeg_channels: int = 32,
                 eeg_seq_len: int = 1280,    # 5s @ 256Hz
                 eye_features: int = 8,
                 driving_features: int = 6,
                 d_model: int = 128,
                 n_heads: int = 4,
                 n_classes: int = 3):  # focused, distracted, drowsy
        super().__init__()
        
        # === EEG 编码器 ===
        self.eeg_encoder = EEGEncoder(
            n_channels=eeg_channels,
            seq_len=eeg_seq_len,
            d_model=d_model
        )
        
        # === 眼动编码器 ===
        self.eye_encoder = EyeMovementEncoder(
            input_dim=eye_features,
            d_model=d_model
        )
        
        # === 驾驶行为编码器 ===
        self.driving_encoder = DrivingBehaviorEncoder(
            input_dim=driving_features,
            d_model=d_model
        )
        
        # === 跨模态注意力融合 ===
        self.cross_modal_attention = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=n_heads,
            batch_first=True
        )
        
        # === 时序建模 ===
        self.temporal_model = nn.LSTM(
            input_size=d_model * 3,
            hidden_size=d_model,
            num_layers=2,
            batch_first=True,
            bidirectional=True
        )
        
        # === 分类头 ===
        self.classifier = nn.Sequential(
            nn.Linear(d_model * 2, d_model),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(d_model, n_classes)
        )
        
        # === 模态权重学习 ===
        self.modal_weights = nn.Parameter(torch.ones(3))  # EEG, Eye, Driving
        
    def forward(self, eeg, eye, driving):
        """
        Args:
            eeg: (B, C, T) EEG信号
            eye: (B, T, F_eye) 眼动特征序列
            driving: (B, T, F_drive) 驾驶行为特征序列
            
        Returns:
            logits: (B, n_classes)
            modal_weights: (3,) 各模态权重
        """
        # 单模态编码
        eeg_feat = self.eeg_encoder(eeg)      # (B, d_model)
        eye_feat = self.eye_encoder(eye)      # (B, d_model)
        drive_feat = self.driving_encoder(driving)  # (B, d_model)
        
        # 模态权重（soft normalization）
        weights = torch.softmax(self.modal_weights, dim=0)
        
        # 加权融合
        fused = (
            weights[0] * eeg_feat + 
            weights[1] * eye_feat + 
            weights[2] * drive_feat
        )
        
        # 跨模态注意力
        modal_stack = torch.stack([eeg_feat, eye_feat, drive_feat], dim=1)  # (B, 3, d_model)
        attn_out, _ = self.cross_modal_attention(
            modal_stack, modal_stack, modal_stack
        )
        fused_attn = attn_out.mean(dim=1)  # (B, d_model)
        
        # 最终融合
        final_feat = torch.cat([fused, fused_attn], dim=-1)
        
        # 分类
        logits = self.classifier(final_feat)
        
        return logits, weights


class EEGEncoder(nn.Module):
    """EEG 编码器：频段功率 + 功能连接性"""
    
    def __init__(self, n_channels=32, seq_len=1280, d_model=128):
        super().__init__()
        self.n_channels = n_channels
        
        # 1D CNN for raw EEG
        self.conv = nn.Sequential(
            nn.Conv1d(n_channels, 64, kernel_size=64, stride=4),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.MaxPool1d(4),
            
            nn.Conv1d(64, 128, kernel_size=32, stride=2),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.MaxPool1d(4),
            
            nn.Conv1d(128, d_model, kernel_size=16, stride=1),
            nn.BatchNorm1d(d_model),
            nn.ReLU(),
            nn.AdaptiveAvgPool1d(1)
        )
        
        # 频段功率提取 (θ: 4-8Hz, α: 8-13Hz, β: 13-30Hz)
        self.band_projection = nn.Linear(d_model + 3 * n_channels, d_model)
        
    def forward(self, x):
        """
        Args:
            x: (B, C, T) raw EEG
        Returns:
            features: (B, d_model)
        """
        # CNN features
        cnn_feat = self.conv(x).squeeze(-1)  # (B, d_model)
        
        # 频段功率（简化模拟）
        # 实际实现需要STFT提取θ/α/β功率
        band_power = torch.randn(x.size(0), 3 * self.n_channels, device=x.device)
        
        # 融合
        combined = torch.cat([cnn_feat, band_power], dim=-1)
        output = self.band_projection(combined)
        
        return output


class EyeMovementEncoder(nn.Module):
    """眼动特征编码器"""
    
    def __init__(self, input_dim=8, d_model=128):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 64),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, d_model)
        )
        
    def forward(self, x):
        """x: (B, T, F) -> mean over T -> (B, d_model)"""
        encoded = self.encoder(x)  # (B, T, d_model)
        return encoded.mean(dim=1)  # (B, d_model)


class DrivingBehaviorEncoder(nn.Module):
    """驾驶行为编码器"""
    
    def __init__(self, input_dim=6, d_model=128):
        super().__init__()
        self.encoder = nn.LSTM(
            input_size=input_dim,
            hidden_size=d_model,
            num_layers=2,
            batch_first=True,
            bidirectional=False
        )
        
    def forward(self, x):
        """x: (B, T, F) -> (B, d_model)"""
        _, (hn, _) = self.encoder(x)
        return hn[-1]  # (B, d_model)


# === 测试代码 ===
if __name__ == "__main__":
    B = 4  # batch size
    
    # 模拟输入
    eeg = torch.randn(B, 32, 1280)      # 32通道, 5秒@256Hz
    eye = torch.randn(B, 50, 8)         # 50帧, 8个眼动特征
    driving = torch.randn(B, 50, 6)      # 50帧, 6个驾驶特征
    
    # 创建模型
    model = MultiModalAttentionDetector()
    
    # 前向传播
    logits, weights = model(eeg, eye, driving)
    
    print(f"EEG输入: {eeg.shape}")
    print(f"眼动输入: {eye.shape}")
    print(f"驾驶输入: {driving.shape}")
    print(f"输出logits: {logits.shape}")
    print(f"模态权重: EEG={weights[0]:.3f}, 眼动={weights[1]:.3f}, 驾驶={weights[2]:.3f}")
    print(f"参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M")
    
    # 模拟性能对比
    print("\n=== 单模态 vs 多模态性能对比 ===")
    print(f"{'模态':<20} {'准确率':>8} {'跨人群泛化':>12}")
    print("-" * 42)
    print(f"{'EEG only':<20} {'72.3%':>8} {'差':>12}")
    print(f"{'眼动 only':<20} {'68.5%':>8} {'中':>12}")
    print(f"{'驾驶行为 only':<20} {'65.1%':>8} {'安全组好/危险组差':>12}")
    print(f"{'EEG+眼动':<20} {'79.8%':>8} {'中':>12}")
    print(f"{'EEG+眼动+驾驶':<20} {'84.2%':>8} {'好':>12}")

# 车载可实现的认知分心检测方案
IN_VEHICLE_FEATURES = {
    'eye_tracking': {
        'features': ['gaze_direction', 'blink_rate', 'blink_duration', 
                     'pupil_diameter', 'saccade_velocity', 'fixation_duration'],
        'sensor': 'IR摄像头 (OV2311)',
        'sampling_rate': '30fps',
        'latency': '<100ms'
    },
    'driving_behavior': {
        'features': ['steering_angle_variance', 'lane_departure_frequency',
                     'speed_variance', 'brake_reaction_time', 
                     'accelerator_pedal_variance', 'heading_error'],
        'sensor': 'CAN总线',
        'sampling_rate': '100Hz',
        'latency': '<50ms'
    },
    'physiological_proxy': {
        'features': ['pupil_diameter_variation', 'heart_rate_variability',
                     'skin_conductance', 'respiration_rate'],
        'sensor': '方向盘传感器+座椅传感器',
        'sampling_rate': '1-10Hz',
        'latency': '<1s'
    }
}

# 替代EEG的方案
print("=== 车载认知分心检测替代EEG方案 ===")
print()
print("| 替代指标        | 对应EEG功能        | 可观测性 | 预测能力 |")
print("|----------------|-------------------|---------|---------|")
print("| 瞳孔直径变化    | 认知负荷 (P3)     | ✅ 高    | 中等    |")
print("| HRV (心率变异)  | 自主神经调控       | ✅ 中    | 中等    |")
print("| 方向盘微颤      | 运动准备 (β)      | ✅ 高    | 中等    |")
print("| 眨眼运动学      | 注意力采样节律     | ✅ 高    | 中等    |")
print("| 车道保持稳定性  | 预备控制 (LFN)    | ✅ 高    | 强      |")
print("| 刹车反应时间    | 刺激评估速度 (P3) | ✅ 中    | 强      |")

class PersonalizedCalibration:
    """个性化DMS校准方案"""
    
    def __init__(self):
        self.baseline = None
        self.calibration_samples = 0
        self.is_calibrated = False
        
    def collect_baseline(self, driver_state):
        """收集正常驾驶基线"""
        if self.baseline is None:
            self.baseline = {}
        
        # 需要收集的基线指标
        metrics = {
            'mean_pupil_diameter': driver_state.pupil_diameter,
            'std_pupil_diameter': driver_state.pupil_diameter_std,
            'mean_blink_duration': driver_state.blink_duration,
            'mean_blink_rate': driver_state.blink_rate,
            'steering_angle_variance': driver_state.steering_angle_var,
            'lane_keeping_error': driver_state.lane_error,
            'brake_reaction_time': driver_state.brake_rt
        }
        
        # 指数移动平均更新基线
        alpha = 0.05  # 慢更新
        for key, value in metrics.items():
            if key not in self.baseline:
                self.baseline[key] = value
            else:
                self.baseline[key] = (1-alpha) * self.baseline[key] + alpha * value
        
        self.calibration_samples += 1
        if self.calibration_samples > 300:  # 10分钟@30fps
            self.is_calibrated = True
    
    def detect_cognitive_distraction(self, current_state):
        """检测认知分心"""
        if not self.is_calibrated:
            return False, "正在校准中..."
        
        # 计算各指标偏离基线的程度
        deviations = {}
        for key, baseline_val in self.baseline.items():
            current_val = getattr(current_state, key, baseline_val)
            if baseline_val > 0:
                deviations[key] = abs(current_val - baseline_val) / baseline_val
        
        # 综合评分
        distraction_score = 0
        weights = {
            'pupil_diameter': 0.2,    # 瞳孔变化反映认知负荷
            'blink_duration': 0.15,   # 眨眼模式变化
            'steering_angle_var': 0.25,  # 方向盘稳定性（最强指标）
            'lane_keeping_error': 0.2,
            'brake_reaction_time': 0.2  # 反应时间
        }
        
        for metric, weight in weights.items():
            for key, dev in deviations.items():
                if metric in key:
                    distraction_score += weight * min(dev, 1.0)
        
        is_distracted = distraction_score > 0.4
        
        return is_distracted, distraction_score