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

class EEGIntentionDetector(nn.Module):
    """
    EEG 驾驶意图检测器
    
    基于 arXiv 2026 论文方法：
    - 原始 EEG 信号输入（无需预处理）
    - 深度学习自动提取特征
    - 三分类：左转/右转/直行
    
    硬件配置：
    | 组件 | 型号 | 参数 | 佩戴方式 |
    |------|------|------|----------|
    | EEG 头带 | Bitbrain Diadem | 16ch, 500Hz | 头带 |
    | 干电极 | g.Sahara | Ag/AgCl 干电极 | 头皮接触 |
    | 处理器 | 车载 GPU | 20W | 中控 |
    | 延迟 | 端到端 | <50ms | - |
    
    电极位置（10-20 系统）：
    - 运动区：C3, Cz, C4（转向想象主要来源）
    - 前额：F3, Fz, F4（执行控制）
    - 顶叶：P3, Pz, P4（空间注意）
    """
    
    def __init__(self, num_channels: int = 16, num_classes: int = 3, 
                 sample_rate: int = 500, window_sec: float = 1.0):
        super().__init__()
        self.num_channels = num_channels
        self.sample_rate = sample_rate
        self.window_samples = int(window_sec * sample_rate)
        
        # 时空双流架构
        # 空间流：通道维度卷积
        self.spatial_conv = nn.Sequential(
            # 输入: (B, 1, num_channels, window_samples)
            nn.Conv2d(1, 32, kernel_size=(num_channels, 1)),  # 跨通道卷积
            nn.BatchNorm2d(32),
            nn.ReLU(),
        )
        
        # 时序流：时间维度卷积
        self.temporal_conv = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=(1, 25), padding=(0, 12)),  # 50ms 窗口
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d((1, 5)),  # 5x 下采样
            
            nn.Conv2d(64, 128, kernel_size=(1, 10), padding=(0, 5)),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d((1, 3)),
        )
        
        # 分类头
        feat_size = 128 * (self.window_samples // 15)  # 估算
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(64, num_classes)
        )
    
    def forward(self, eeg: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        Args:
            eeg: (B, num_channels, window_samples) 原始 EEG
            
        Returns:
            output: {
                logits: (B, 3) 分类 logits,
                confidence: (B,) 最大概率,
                prediction: (B,) 预测类别
            }
        """
        # 添加通道维度
        x = eeg.unsqueeze(1)  # (B, 1, C, T)
        
        # 空间卷积
        x = self.spatial_conv(x)  # (B, 32, 1, T)
        
        # 时序卷积
        x = self.temporal_conv(x)  # (B, 128, 1, T')
        
        # 分类
        logits = self.classifier(x)
        
        probs = torch.softmax(logits, dim=-1)
        confidence, prediction = probs.max(dim=-1)
        
        return {
            "logits": logits,
            "confidence": confidence,
            "prediction": prediction,
            "probabilities": probs,
        }


# 测试
if __name__ == "__main__":
    detector = EEGIntentionDetector(
        num_channels=16, 
        num_classes=3,
        sample_rate=500,
        window_sec=1.0
    )
    
    # 模拟 EEG 输入
    eeg = torch.randn(4, 16, 500)  # 4个样本, 16通道, 1秒
    
    with torch.no_grad():
        result = detector(eeg)
    
    print("=== EEG 意图检测 ===")
    classes = ["left_turn", "right_turn", "straight"]
    for i, cls in enumerate(classes):
        print(f"  {cls}: {result['probabilities'][0, i]:.4f}")
    print(f"预测: {classes[result['prediction'][0]]}")
    print(f"置信度: {result['confidence'][0]:.4f}")
    print(f"\n参数量: {sum(p.numel() for p in detector.parameters()):,}")

class BCIApplicationsInCabin:
    """
    BCI 在汽车座舱的三层应用架构
    
    L1: 驾驶意图预测（B2V）
       - 转向意图提前 0.5s 检测
       - 刹车意图提前 0.3s 检测
       - ADAS 提前介入
    
    L2: 认知状态监测
       - 认知负荷评估（theta/alpha 比值）
       - 疲劳检测（alpha 波增加）
       - 注意力分散（P300 振幅降低）
    
    L3: 无形交互
       - 意念控制信息娱乐
       - 意念选择菜单
       - 错误电位反馈（系统误解时自动纠正）
    """
    
    APPLICATIONS = {
        "L1_intention": {
            "name": "驾驶意图预测",
            "paradigm": "Motor Imagery (MI)",
            "latency": "0.5s 提前",
            "accuracy": "80%+",
            "use_case": "ADAS 提前介入",
            "safety_level": "安全关键",
        },
        "L1_braking": {
            "name": "刹车意图检测",
            "paradigm": "ErrP + MI 混合",
            "latency": "0.3s 提前",
            "accuracy": "85%+",
            "use_case": "紧急刹车辅助",
            "safety_level": "安全关键",
        },
        "L2_cognitive": {
            "name": "认知负荷评估",
            "paradigm": "频谱分析",
            "latency": "30-60s 累积",
            "accuracy": "75%+",
            "use_case": "DMS 增强输入",
            "safety_level": "辅助",
        },
        "L2_fatigue": {
            "name": "疲劳检测",
            "paradigm": "Alpha 功率增加",
            "latency": "60s 累积",
            "accuracy": "80%+",
            "use_case": "DMS 疲劳增强",
            "safety_level": "辅助",
        },
        "L3_interaction": {
            "name": "意念菜单选择",
            "paradigm": "P300 / SSVEP",
            "latency": "1-3s",
            "accuracy": "85-95%",
            "use_case": "信息娱乐控制",
            "safety_level": "便利",
        },
        "L3_error": {
            "name": "错误电位反馈",
            "paradigm": "ErrP",
            "latency": "0.3-0.5s",
            "accuracy": "70%+",
            "use_case": "系统误识别自动纠正",
            "safety_level": "便利",
        },
    }