import torch
import torch.nn as nn
import numpy as np
from scipy import signal as sig

class UnifiedEEGFeatureExtractor(nn.Module):
    """
    统一EEG特征提取器
    
    目标：提取跨被试通用的EEG特征
    关键：使用域自适应技术，减少个体差异
    
    参考: Ammar et al., Electronics 2026
    """
    
    def __init__(self, num_channels=14, fs=128, num_classes=3):
        super().__init__()
        self.num_channels = num_channels
        self.fs = fs
        
        # 1. 时空特征提取
        self.spatial_filter = nn.Conv2d(1, 32, (num_channels, 1))
        self.temporal_filter = nn.Conv2d(32, 64, (1, 25), padding=(0, 12))
        
        # 2. 域自适应层（减少跨被试差异）
        self.domain_adapt = nn.Sequential(
            nn.Linear(64 * 128, 256),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(256, 128),
            nn.ReLU()
        )
        
        # 3. 分类器
        self.classifier = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, num_classes)  # alert/drowsy/distracted
        )
    
    def forward(self, x):
        """
        Args:
            x: (B, 1, 14, 128) - 14通道, 1秒@128Hz
        
        Returns:
            logits: (B, 3)
        """
        x = self.spatial_filter(x)    # (B, 32, 1, T)
        x = torch.relu(x)
        x = self.temporal_filter(x)   # (B, 64, 1, T)
        x = torch.relu(x)
        x = x.flatten(1)              # (B, 64*T)
        x = self.domain_adapt(x)      # (B, 128)
        return self.classifier(x)     # (B, 3)


class LeakageFreePipeline:
    """
    防数据泄露的实验Pipeline
    
    论文核心贡献之一：严格的数据隔离协议
    """
    
    def __init__(self):
        self.train_subjects = []
        self.test_subjects = []
    
    def leave_one_subject_out(self, all_subjects):
        """
        留一被试交叉验证（LOSOCV）
        
        关键：训练集中完全不包含测试被试的任何数据
        包括：不同时段、不同会话的数据
        """
        results = []
        for test_subj in all_subjects:
            train_subs = [s for s in all_subjects if s != test_subj]
            
            # 严格隔离
            train_data = self.load_data(train_subs)
            test_data = self.load_data([test_subj])
            
            # 训练
            model = UnifiedEEGFeatureExtractor()
            self.train(model, train_data)
            
            # 测试
            acc = self.evaluate(model, test_data)
            results.append((test_subj, acc))
        
        return results
    
    def check_leakage(self, train_data, test_data):
        """检查数据泄露"""
        # 1. 确保无被试重叠
        train_subj = set(train_data['subject'])
        test_subj = set(test_data['subject'])
        assert len(train_subj & test_subj) == 0, "数据泄露：训练/测试有被试重叠"
        
        # 2. 确保无时间泄露
        # 测试数据时间戳必须完全独立
        
        # 3. 确保无特征泄露
        # 不能使用测试集统计量做归一化
        return True

class SafeStopTrigger:
    """
    L3安全停止触发机制
    
    基于EEG驾驶员警觉性评估
    当警觉性低于阈值时触发MRC（最小风险策略）
    """
    
    VIGILANCE_THRESHOLD = 0.65
    CONFIRMATION_WINDOW = 5  # 连续5秒低于阈值才触发
    TAKEOVER_REQUEST_TIME = 4  # TOR后给4秒响应时间
    
    def __init__(self, eeg_model):
        self.model = eeg_model
        self.vigilance_history = []
        self.state = "L3_ACTIVE"
    
    def monitor_vigilance(self, eeg_stream, fs=128):
        """
        实时监控驾驶员警觉性
        
        状态机：
        L3_ACTIVE → TOR_ISSUED → DRIVER_RESPONDED → L2_HANDOVER
                  → TOR_ISSUED → NO_RESPONSE → MRC_SAFE_STOP
        """
        window = eeg_stream  # 1秒窗口
        
        with torch.no_grad():
            x = torch.FloatTensor(window).unsqueeze(0).unsqueeze(0)
            logits = self.model(x)
            probs = torch.softmax(logits, dim=1)
        
        vigilance = probs[0, 0].item()  # alert概率
        self.vigilance_history.append(vigilance)
        
        # 保持最近5秒
        if len(self.vigilance_history) > self.CONFIRMATION_WINDOW:
            self.vigilance_history.pop(0)
        
        # 判断是否需要触发
        if len(self.vigilance_history) == self.CONFIRMATION_WINDOW:
            avg_vigilance = np.mean(self.vigilance_history)
            
            if avg_vigilance < self.VIGILANCE_THRESHOLD:
                return self.trigger_safe_stop()
        
        return "MONITORING", avg_vigilance
    
    def trigger_safe_stop(self):
        """触发安全停止流程"""
        if self.state == "L3_ACTIVE":
            self.state = "TOR_ISSUED"
            return "TAKEOVER_REQUEST", None
        
        elif self.state == "TOR_ISSUED":
            # 检查驾驶员是否响应
            if not self.driver_responded():
                self.state = "MRC_EXECUTING"
                return "MINIMUM_RISK_MANEUVER", None
            else:
                self.state = "L2_HANDOVER"
                return "DRIVER_TAKEOVER", None
        
        return "MONITORING", None
    
    def driver_responded(self):
        """检测驾驶员是否响应TOR"""
        # 检查方向盘输入、踏板输入、头部运动
        # 简化示意
        return False

class DomainAdaptiveTraining:
    """
    域自适应训练策略
    
    解决EEG跨被试泛化问题
    """
    
    def __init__(self):
        self.feature_extractor = UnifiedEEGFeatureExtractor()
        self.domain_classifier = nn.Sequential(
            nn.Linear(128, 32),
            nn.ReLU(),
            nn.Linear(32, num_subjects)  # 被试身份分类
        )
    
    def adversarial_loss(self, features, subject_labels):
        """
        对抗损失：让特征无法区分被试身份
        
        目标：特征提取器学习"被试无关"的EEG特征
        """
        domain_pred = self.domain_classifier(features)
        # 梯度反转层：最大化domain loss
        return -nn.CrossEntropyLoss()(domain_pred, subject_labels)
    
    def train_step(self, eeg_data, state_labels, subject_labels):
        """
        多任务训练：
        1. 状态分类（主任务）
        2. 被试身份分类（对抗任务，梯度反转）
        """
        features = self.feature_extractor(eeg_data)
        
        # 主任务：驾驶员状态
        state_pred = self.classifier(features)
        state_loss = nn.CrossEntropyLoss()(state_pred, state_labels)
        
        # 对抗任务：被试身份
        domain_loss = self.adversarial_loss(features, subject_labels)
        
        total_loss = state_loss + 0.1 * domain_loss
        return total_loss