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

class EEGSteeringIntentionCNN(nn.Module):
    """
    EEG转向意图预测CNN
    
    论文: arXiv:2601.05084
    输入: 原始EEG信号 (14通道, 1秒窗口)
    输出: 三分类 (左转/右转/直行)
    准确率: 83.7%
    
    关键创新: 无需预处理，CNN直接学习原始EEG
    """
    
    def __init__(self, n_channels=14, n_classes=3, 
                 sample_rate=128, window_s=1):
        super().__init__()
        
        self.n_channels = n_channels
        self.window_samples = sample_rate * window_s
        
        # 第一层：时间卷积
        self.temporal_conv = nn.Sequential(
            nn.Conv2d(1, 25, kernel_size=(1, 10), stride=(1, 1)),
            nn.BatchNorm2d(25),
            nn.ELU()
        )
        
        # 第二层：空间卷积（跨通道）
        self.spatial_conv = nn.Sequential(
            nn.Conv2d(25, 50, kernel_size=(n_channels, 1), stride=(1, 1)),
            nn.BatchNorm2d(50),
            nn.ELU(),
            nn.MaxPool2d(kernel_size=(1, 3), stride=(1, 3))
        )
        
        # 第三层：特征提取
        self.feature_conv = nn.Sequential(
            nn.Conv2d(50, 100, kernel_size=(1, 10), stride=(1, 1)),
            nn.BatchNorm2d(100),
            nn.ELU(),
            nn.MaxPool2d(kernel_size=(1, 3), stride=(1, 3))
        )
        
        # 计算展平后维度
        with torch.no_grad():
            dummy = torch.zeros(1, 1, n_channels, self.window_samples)
            dummy = self.temporal_conv(dummy)
            dummy = self.spatial_conv(dummy)
            dummy = self.feature_conv(dummy)
            flat_size = dummy.numel()
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Dropout(0.5),
            nn.Linear(flat_size, 128),
            nn.ELU(),
            nn.Dropout(0.5),
            nn.Linear(128, n_classes)
        )
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: 原始EEG, shape=(B, 1, n_channels, window_samples)
            
        Returns:
            logits: shape=(B, 3) [left, straight, right]
        """
        x = self.temporal_conv(x)
        x = self.spatial_conv(x)
        x = self.feature_conv(x)
        logits = self.classifier(x)
        return logits


class B2VPipeline:
    """
    Brain-to-Vehicle通信管道
    
    论文验证: 模拟器环境
    准确率: 83.7% (三分类)
    延迟: CNN推理 <10ms
    
    关键发现:
    1. 右转分类最准（大脑空间偏侧性）
    2. 无需预处理（CNN直接学习原始信号）
    3. 比传统EEG方法高16-34%
    """
    
    INTENTIONS = ['left_turn', 'straight', 'right_turn']
    
    def __init__(self):
        self.model = EEGSteeringIntentionCNN()
        self.model.eval()
        
        # 置信度阈值
        self.CONFIDENCE_THRESHOLD = 0.6
        
    def predict_intention(self, eeg_window: np.ndarray) -> dict:
        """
        实时预测转向意图
        
        Args:
            eeg_window: 1秒EEG窗口, shape=(n_channels, n_samples)
            
        Returns:
            dict: 预测结果
        """
        with torch.no_grad():
            x = torch.FloatTensor(eeg_window).unsqueeze(0).unsqueeze(0)
            logits = self.model(x)
            probs = torch.softmax(logits, dim=1)
            
            pred_idx = probs.argmax(1).item()
            confidence = probs[0, pred_idx].item()
            
            return {
                'intention': self.INTENTIONS[pred_idx],
                'confidence': confidence,
                'probabilities': {
                    intent: probs[0, i].item()
                    for i, intent in enumerate(self.INTENTIONS)
                },
                'ready_for_b2v': confidence > self.CONFIDENCE_THRESHOLD
            }


# 实际测试
if __name__ == "__main__":
    print("=== B2V 转向意图预测 ===\n")
    
    pipeline = B2VPipeline()
    
    # 模拟1秒EEG数据 (14通道, 128Hz)
    np.random.seed(42)
    n_channels, n_samples = 14, 128
    eeg = np.random.normal(0, 10, (n_channels, n_samples))
    
    result = pipeline.predict_intention(eeg)
    
    print(f"预测意图: {result['intention']}")
    print(f"置信度: {result['confidence']:.2%}")
    print(f"B2V就绪: {result['ready_for_b2v']}")
    print(f"\n各类别概率:")
    for intent, prob in result['probabilities'].items():
        print(f"  {intent}: {prob:.2%}")
    
    print(f"\n=== 论文结果 ===")
    print(f"  三分类准确率: 83.7%")
    print(f"  右转准确率最高 (空间偏侧性)")
    print(f"  传统方法: 50-67% (三分类)")
    print(f"  提升: +16-34%")
    print(f"  预处理: 无需（CNN直接学习原始EEG）")
    print(f"  延迟: CNN推理 <10ms")
    print(f"  环境: 模拟器（未在真实驾驶验证）")

class DriverIntentionPredictor:
    """
    多模态驾驶员意图预测
    
    B2V(EEG)太远，但可以融合其他信号
    提前预测驾驶员行为
    """
    
    def __init__(self):
        # 可用信号源
        self.signals = {
            'gaze_trajectory': {'weight': 0.30, 'lead_time_ms': 500},
            'head_orientation': {'weight': 0.20, 'lead_time_ms': 400},
            'foot_position': {'weight': 0.15, 'lead_time_ms': 300},
            'hand_grip': {'weight': 0.10, 'lead_time_ms': 200},
            'steering_torque': {'weight': 0.15, 'lead_time_ms': 100},
            'pedal_pressure': {'weight': 0.10, 'lead_time_ms': 100},
        }
        
    def predict_lane_change(self, signals: dict) -> dict:
        """
        预测变道意图
        
        视线是最早的意图信号（提前500ms）
        """
        score = 0
        for signal, value in signals.items():
            w = self.signals.get(signal, {}).get('weight', 0)
            score += w * value
        
        if score > 0.7:
            return {'intention': 'lane_change', 'confidence': score, 'lead_time_ms': 500}
        elif score > 0.4:
            return {'intention': 'maybe_lane_change', 'confidence': score, 'lead_time_ms': 300}
        else:
            return {'intention': 'stay_in_lane', 'confidence': 1-score, 'lead_time_ms': 0}