"""
EEG + 眼动融合认知状态监控
基于 Frontiers in Neuroergonomics (2025) 系统综述

融合策略对比：
1. 早期融合（特征拼接）: 简单但维度灾难
2. 中期融合（各自特征提取后拼接）: 推荐
3. 晚期融合（各自分类后投票）: 灵活但信息损失

综述发现：中期融合在准确率和鲁棒性间最佳平衡
"""

import numpy as np
import torch
import torch.nn as nn
from typing import Tuple, List, Optional

class EEGFeatureExtractor(nn.Module):
    """
    EEG 特征提取器
    
    提取特征：
    1. 频带功率：δ(1-4Hz), θ(4-7Hz), α(8-12Hz), β(13-30Hz), γ(30-80Hz)
    2. 频带功率比：θ/α（疲劳指标）, β/α（警觉指标）
    3. 事件相关电位：P300振幅和潜伏期
    4. 连接性：相位锁定值(PLV)
    
    综述关键发现：
    - θ波增加 + α波减少 → 疲劳
    - β波减少 → 警觉性下降
    - P300潜伏期延长 → 认知负荷增加
    """
    
    def __init__(self, num_channels=64, sampling_rate=512):
        super().__init__()
        self.num_channels = num_channels
        self.sampling_rate = sampling_rate
        
        # 1D-CNN 提取时域特征
        self.temporal_conv = nn.Sequential(
            nn.Conv1d(num_channels, 128, 64, stride=4, padding=30),
            nn.BatchNorm1d(128),
            nn.ELU(),
            nn.MaxPool1d(4),
            
            nn.Conv1d(128, 256, 16, stride=2, padding=7),
            nn.BatchNorm1d(256),
            nn.ELU(),
            nn.MaxPool1d(4),
            
            nn.AdaptiveAvgPool1d(1),
        )
        
        # 输出维度
        self.feature_dim = 256
    
    def extract_band_power(self, eeg_signal: np.ndarray) -> dict:
        """
        提取频带功率
        
        Args:
            eeg_signal: (num_channels, T) EEG信号
        
        Returns:
            band_powers: 各频带功率
        """
        from scipy.signal import welch
        
        freqs, psd = welch(eeg_signal, fs=self.sampling_rate, nperseg=1024)
        
        bands = {
            'delta': (1, 4),
            'theta': (4, 7),
            'alpha': (8, 12),
            'beta': (13, 30),
            'gamma': (30, 80),
        }
        
        powers = {}
        for name, (low, high) in bands.items():
            mask = (freqs >= low) & (freqs < high)
            powers[name] = np.mean(psd[:, mask], axis=1) if len(psd.shape) > 1 else np.mean(psd[mask])
        
        # 比率指标
        powers['theta_alpha_ratio'] = np.mean(powers['theta']) / (np.mean(powers['alpha']) + 1e-8)
        powers['beta_alpha_ratio'] = np.mean(powers['beta']) / (np.mean(powers['alpha']) + 1e-8)
        
        return powers
    
    def forward(self, x):
        """x: (B, C, T)"""
        features = self.temporal_conv(x)
        return features.view(x.size(0), -1)  # (B, 256)


class EyeTrackingFeatureExtractor(nn.Module):
    """
    眼动特征提取器
    
    提取特征：
    1. 注视参数：注视点位置、持续时间、注视熵
    2. 扫视参数：扫视频率、幅度、方向
    3. 眨眼参数：眨眼频率、时长、PERCLOS
    4. 瞳孔参数：直径、变化率、瞳孔 unrest
    
    综述关键发现：
    - PERCLOS > 80% → 微睡眠
    - 注视平坦（0.5-15秒）+ 闭眼 → 微睡眠确认
    - 瞳孔直径下降 → 疲劳
    - 扫视频率降低 → 警觉性下降
    - 注视熵增加 → 认知负荷增加
    """
    
    def __init__(self, input_dim=6):
        """
        输入: (gaze_x, gaze_y, pupil_diameter, eye_openness, blink_flag, timestamp)
        """
        super().__init__()
        
        # LSTM 提取时序特征
        self.lstm = nn.LSTM(
            input_size=input_dim,
            hidden_size=128,
            num_layers=2,
            batch_first=True,
            bidirectional=True,
            dropout=0.2,
        )
        
        # 注意力层
        self.attention = nn.Sequential(
            nn.Linear(256, 64),
            nn.Tanh(),
            nn.Linear(64, 1),
            nn.Softmax(dim=1),
        )
        
        self.feature_dim = 256
    
    def forward(self, x):
        """x: (B, T, 6)"""
        lstm_out, _ = self.lstm(x)  # (B, T, 256)
        
        # Attention
        att_weights = self.attention(lstm_out)  # (B, T, 1)
        context = torch.sum(att_weights * lstm_out, dim=1)  # (B, 256)
        
        return context


class MultimodalFusion(nn.Module):
    """
    多模态融合分类器
    
    中期融合策略（综述推荐）：
    1. EEG → CNN → 256维特征
    2. 眼动 → LSTM+Attention → 256维特征
    3. 拼接 → 512维
    4. MLP → 分类
    
    分类输出：
    - 0: 警觉 (alert)
    - 1: 轻度疲劳 (mild drowsiness)
    - 2: 微睡眠 (microsleep)
    - 3: 认知负荷高 (high cognitive load)
    - 4: 视觉分心 (visual distraction)
    - 5: 认知分心 (cognitive distraction)
    """
    
    def __init__(self, num_classes=6):
        super().__init__()
        
        self.eeg_extractor = EEGFeatureExtractor()
        self.eye_extractor = EyeTrackingFeatureExtractor()
        
        # 融合层
        self.fusion = nn.Sequential(
            nn.Linear(self.eeg_extractor.feature_dim + 
                      self.eye_extractor.feature_dim, 256),
            nn.BatchNorm1d(256),
            nn.ELU(),
            nn.Dropout(0.3),
            
            nn.Linear(256, 128),
            nn.BatchNorm1d(128),
            nn.ELU(),
            nn.Dropout(0.2),
            
            nn.Linear(128, num_classes),
        )
    
    def forward(self, eeg_input, eye_input):
        """
        Args:
            eeg_input: (B, 64, T) EEG信号
            eye_input: (B, T, 6) 眼动参数序列
        """
        eeg_features = self.eeg_extractor(eeg_input)  # (B, 256)
        eye_features = self.eye_extractor(eye_input)  # (B, 256)
        
        # 拼接
        fused = torch.cat([eeg_features, eye_features], dim=1)  # (B, 512)
        
        # 分类
        output = self.fusion(fused)  # (B, 6)
        
        return output


# ==================== 测试 ====================
if __name__ == "__main__":
    np.random.seed(42)
    
    print("=== EEG + 眼动融合认知状态监控 ===")
    
    model = MultimodalFusion(num_classes=6)
    
    # 模拟输入
    batch_size = 4
    eeg_seq_len = 1024  # 2秒@512Hz
    eye_seq_len = 60    # 2秒@30Hz
    
    eeg_input = torch.randn(batch_size, 64, eeg_seq_len)
    eye_input = torch.randn(batch_size, eye_seq_len, 6)
    
    # 前向传播
    output = model(eeg_input, eye_input)
    
    print(f"EEG输入: {eeg_input.shape}")
    print(f"眼动输入: {eye_input.shape}")
    print(f"输出: {output.shape}")
    print(f"预测类别: {torch.argmax(output, dim=1)}")
    
    # 模型参数
    total_params = sum(p.numel() for p in model.parameters())
    print(f"\n模型参数量: {total_params:,}")
    print(f"模型大小: ~{total_params * 4 / 1024 / 1024:.1f} MB")
    
    # 性能对比
    print(f"\n=== 融合 vs 单模态性能对比 ===")
    print(f"{'状态':<15} {'EEG':<10} {'眼动':<10} {'融合':<10} {'提升':<10}")
    states = [
        ('警觉/疲劳', 93, 91, 97.5),
        ('微睡眠', 95, 96, 97.5),
        ('认知负荷', 87, 72, 89),
        ('情绪状态', 78, 65, 82),
        ('分心类型', 70, 85, 90),
    ]
    for name, eeg, eye, fused in states:
        improvement = fused - max(eeg, eye)
        print(f"{name:<15} {eeg}%{'':<5} {eye}%{'':<5} {fused}%{'':<5} +{improvement:.1f}%")
    
    # 互补性分析
    print(f"\n=== 互补性分析 ===")
    print(f"EEG 优势: 毫秒级时间分辨率、直接神经信号、微睡眠早期")
    print(f"眼动优势: 行为表现指标、非接触采集、PERCLOS标准化")
    print(f"EEG 劣势: 穿戴不便、运动伪影、个体差异大")
    print(f"眼动劣势: 间接信号、行为出现后才检测、光照敏感")
    print(f"融合互补: EEG早期检测 + 眼动行为确认 → 减少+5%误报")
    
    # 微睡眠定义共识
    print(f"\n=== 微睡眠定义共识 (Frontiers 2025) ===")
    print(f"定义: 0.5-15秒闭眼 + 注视平坦/不连贯 + PERCLOS>80%")
    print(f"EEG特征: θ波增加(4-7Hz) + α波减少(8-12Hz)")
    print(f"眼动特征: 注视丢失 + 眨眼时长>500ms + PERCLOS>80%")
    print(f"融合判定: 两者同时出现 → 确认微睡眠(97.5%)")
    
    # 部署建议
    print(f"\n=== 部署建议 ===")
    print(f"Phase 1: 纯眼动DMS（已量产） → 88-96%准确率")
    print(f"Phase 2: 眼动+EEG可选融合 → 97.5%准确率")
    print(f"  - EEG作为高端选配（干电极头带/座椅头枕）")
    print(f"  - 非接触眼动为主，EEG增强关键场景")
    print(f"Phase 3: 全融合（含干电极EEG帽） → 研究/航空")