# 示意代码，不可直接运行（非论文官方实现；论文代码"发表后提供"，尚未有公开仓库。
# 论文的TR-MVPA按时间点逐点训练分类器并以AUC曲线报告结果，特征/分类器细节以原文为准）
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_val_score

class TimeResolvedMVPA:
    """
    时间分辨多变量模式分析
    
    论文核心方法：在毫秒级时间精度上解码认知负荷状态
    使用滑动窗口对每个时间点训练分类器
    """
    
    def __init__(self, n_timepoints: int = 500, 
                 window_ms: int = 100,
                 step_ms: int = 10):
        """
        Args:
            n_timepoints: 总时间点数（采样率×时长）
            window_ms: 解码窗口长度（毫秒）
            step_ms: 步长（毫秒）
        """
        self.window_samples = int(window_ms * 500 / 1000)  # 转换为样本数
        self.step_samples = int(step_ms * 500 / 1000)
        self.classifiers = []
        self.time_centers = []
        
    def fit(self, eeg_data: np.ndarray, labels: np.ndarray):
        """
        训练时序解码器
        
        Args:
            eeg_data: [n_trials, n_timepoints] EEG信号
            labels: [n_trials] 0=低负荷, 1=高负荷
            
        Returns:
            self
        """
        n_trials, n_timepoints = eeg_data.shape
        
        for t in range(0, n_timepoints - self.window_samples, 
                       self.step_samples):
            # 提取时间窗口特征
            window = eeg_data[:, t:t+self.window_samples]
            
            # 提取频域特征
            features = self._extract_features(window)
            
            # 训练分类器
            clf = LogisticRegression(
                C=1.0, penalty='l2', solver='lbfgs',
                max_iter=1000
            )
            clf.fit(features, labels)
            
            self.classifiers.append(clf)
            self.time_centers.append(t + self.window_samples // 2)
            
        return self
    
    def _extract_features(self, window: np.ndarray) -> np.ndarray:
        """
        从EEG窗口提取特征
        
        Args:
            window: [n_trials, n_samples] 
            
        Returns:
            features: [n_trials, n_features]
        """
        # 时域特征
        mean = np.mean(window, axis=1, keepdims=True)
        std = np.std(window, axis=1, keepdims=True)
        skew = self._skewness(window)
        
        # 频域特征（FFT）
        fft = np.fft.rfft(window, axis=1)
        power = np.abs(fft) ** 2
        
        # 关键频段功率
        freqs = np.fft.rfftfreq(window.shape[1], d=1/500)
        
        # Theta频段 (4-8 Hz) - 认知负荷增加时增强
        theta_mask = (freqs >= 4) & (freqs <= 8)
        theta_power = np.sum(power[:, theta_mask], axis=1, keepdims=True)
        
        # Alpha频段 (8-13 Hz) - 认知负荷增加时减弱
        alpha_mask = (freqs >= 8) & (freqs <= 13)
        alpha_power = np.sum(power[:, alpha_mask], axis=1, keepdims=True)
        
        # Beta频段 (13-30 Hz) - 注意集中时增强
        beta_mask = (freqs >= 13) & (freqs <= 30)
        beta_power = np.sum(power[:, beta_mask], axis=1, keepdims=True)
        
        # Theta/Alpha比值（经典认知负荷指标）
        theta_alpha_ratio = theta_power / (alpha_power + 1e-8)
        
        features = np.hstack([
            mean, std, skew,
            theta_power, alpha_power, beta_power,
            theta_alpha_ratio
        ])
        
        return features
    
    def _skewness(self, x: np.ndarray) -> np.ndarray:
        """计算偏度"""
        mean = np.mean(x, axis=1, keepdims=True)
        std = np.std(x, axis=1, keepdims=True)
        return np.mean(((x - mean) / (std + 1e-8)) ** 3, 
                      axis=1, keepdims=True)
    
    def decode_accuracy(self, eeg_data: np.ndarray, 
                        labels: np.ndarray) -> list:
        """
        计算每个时间点的解码准确率
        
        Returns:
            accuracies: 每个时间点的解码准确率列表
        """
        accuracies = []
        for i, clf in enumerate(self.classifiers):
            t = i * self.step_samples
            window = eeg_data[:, t:t+self.window_samples]
            features = self._extract_features(window)
            acc = cross_val_score(clf, features, labels, 
                                  cv=5, scoring='accuracy').mean()
            accuracies.append(acc)
        return accuracies


# 实际测试
if __name__ == "__main__":
    # 模拟EEG数据
    np.random.seed(42)
    n_trials = 100
    n_timepoints = 500 * 30  # 30秒@500Hz
    
    # 低负荷：较强alpha波
    low_load = np.random.randn(n_trials//2, n_timepoints) * 0.5
    low_load += 2.0 * np.sin(2 * np.pi * 10 * 
                               np.arange(n_timepoints) / 500)
    
    # 高负荷：theta增强, alpha减弱
    high_load = np.random.randn(n_trials//2, n_timepoints) * 0.7
    high_load += 3.0 * np.sin(2 * np.pi * 6 * 
                                np.arange(n_timepoints) / 500)
    
    eeg_data = np.vstack([low_load, high_load])
    labels = np.array([0]*(n_trials//2) + [1]*(n_trials//2))
    
    # 训练解码器
    decoder = TimeResolvedMVPA(
        n_timepoints=n_timepoints,
        window_ms=100, step_ms=50
    )
    decoder.fit(eeg_data, labels)
    
    # 解码准确率
    accuracies = decoder.decode_accuracy(eeg_data, labels)
    
    print(f"时间点数: {len(accuracies)}")
    print(f"峰值准确率: {max(accuracies):.2%}")
    print(f"平均准确率: {np.mean(accuracies):.2%}")
    print(f"检测延迟(首次>55%): {next((i*50 for i, a in enumerate(accuracies) if a > 0.55), 'N/A')}ms")

# 示意代码，不可直接运行（多模态融合权重与阈值为本站示意设计，非论文内容；
# 论文未做多模态加权融合，只做了各模态独立解码对比）
class CognitiveDistractionDetector:
    """
    IMS认知分心多模态检测模块
    
    融合：耳道EEG + 眼动 + 头部姿态
    """
    
    def __init__(self):
        # EEG解码器
        self.eeg_decoder = TimeResolvedMVPA(
            window_ms=100, step_ms=50
        )
        # 眼动分析器
        self.gaze_analyzer = GazePatternAnalyzer()
        # 头部姿态分析器
        self.head_analyzer = HeadPoseAnalyzer()
        
        # 多模态融合权重
        self.weights = {
            'eeg': 0.4,      # 神经源头
            'gaze': 0.45,     # 行为后果（最早检测）
            'head': 0.15      # 辅助信号
        }
        
    def detect_cognitive_distraction(
        self, eeg_signal, gaze_data, head_pose
    ) -> dict:
        """
        多模态认知分心检测
        
        Args:
            eeg_signal: [T] 单通道EEG信号
            gaze_data: [T, 2] 视线方向序列
            head_pose: [T, 3] 头部姿态欧拉角
            
        Returns:
            result: {
                'score': 0-1分心概率,
                'modality_scores': 各模态分心概率,
                'latency_ms': 检测延迟
            }
        """
        # 1. EEG解码
        eeg_score = self.eeg_decoder.predict_proba(eeg_signal)
        
        # 2. 眼动模式分析
        gaze_metrics = self.gaze_analyzer.analyze(gaze_data)
        # 关键指标：注视集中度、扫视频率、瞳孔直径
        gaze_score = self._gaze_to_distraction(gaze_metrics)
        
        # 3. 头部姿态分析
        head_score = self.head_analyzer.assess(head_pose)
        
        # 4. 加权融合
        total_score = (
            self.weights['eeg'] * eeg_score +
            self.weights['gaze'] * gaze_score +
            self.weights['head'] * head_score
        )
        
        return {
            'score': total_score,
            'modality_scores': {
                'eeg': eeg_score,
                'gaze': gaze_score,
                'head': head_score
            },
            'latency_ms': 250  # EEG限制
        }
    
    def _gaze_to_distraction(self, metrics):
        """
        眼动指标→认知分心概率
        
        关键指标：
        - 注视集中度：高负荷时视线集中于一点
        - 扫视频率：高负荷时减少（隧道视觉）
        - 瞳孔直径：高负荷时瞳孔扩大
        """
        concentration = metrics['concentration_score']
        saccade_rate = metrics['saccade_rate']
        pupil_size = metrics['pupil_diameter']
        
        # 隧道视觉检测
        if concentration > 0.8 and saccade_rate < 0.3:
            return 0.7
        # 瞳孔扩大+注视集中
        if pupil_size > 1.2 and concentration > 0.6:
            return 0.6
        return 0.3