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

class EEGEyeFusionModel(nn.Module):
    """
    EEG + 眼动跟踪多模态融合模型
    
    基于综述发现的最佳实践：
    - 特征级融合 + 注意力机制
    - EEG: 频域特征 (5频段功率)
    - ET: 行为特征 (注视+眨眼+瞳孔)
    
    论文综述表明注意力融合策略达到最高准确率
    """
    
    def __init__(
        self,
        eeg_channels: int = 14,      # EEG 通道数
        eeg_bands: int = 5,           # 频段数 (δ,θ,α,β,γ)
        et_features: int = 6,         # 眼动特征数
        hidden_dim: int = 128,
        num_classes: int = 3,         # 正常/疲劳/分心
        dropout: float = 0.3
    ):
        super().__init__()
        
        # EEG 特征编码器
        self.eeg_encoder = nn.Sequential(
            nn.Linear(eeg_channels * eeg_bands, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout)
        )
        
        # 眼动特征编码器
        self.et_encoder = nn.Sequential(
            nn.Linear(et_features, hidden_dim // 2),
            nn.BatchNorm1d(hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout)
        )
        
        # 跨模态注意力
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=hidden_dim,
            num_heads=4,
            dropout=dropout,
            batch_first=True
        )
        
        # 融合分类器
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, num_classes)
        )
        
    def forward(
        self,
        eeg_features: torch.Tensor,  # (B, C, B_bands) 通道×频段功率
        et_features: torch.Tensor    # (B, T, et_features) 时序眼动特征
    ) -> torch.Tensor:
        """
        前向传播
        
        Args:
            eeg_features: EEG 频域特征 (批次, 通道, 频段)
            et_features: 眼动时序特征 (批次, 时间步, 特征)
            
        Returns:
            logits: (B, num_classes) 状态分类
        """
        B = eeg_features.shape[0]
        
        # 1. EEG 编码
        eeg_flat = eeg_features.reshape(B, -1)
        eeg_embed = self.eeg_encoder(eeg_flat)  # (B, D)
        
        # 2. 眼动编码（取时序平均）
        et_mean = et_features.mean(dim=1)  # (B, et_features)
        et_embed = self.et_encoder(et_mean)  # (B, D)
        
        # 3. 跨模态注意力
        # EEG 作为 query，ET 作为 key/value
        eeg_q = eeg_embed.unsqueeze(1)  # (B, 1, D)
        et_kv = et_embed.unsqueeze(1)   # (B, 1, D)
        
        attended, _ = self.cross_attn(eeg_q, et_kv, et_kv)
        attended = attended.squeeze(1)  # (B, D)
        
        # 4. 融合分类
        fused = torch.cat([attended, et_embed], dim=-1)  # (B, 2D)
        logits = self.classifier(fused)
        
        return logits


class EEGFeatureExtractor:
    """
    EEG 频域特征提取
    
    将原始 EEG 信号转换为 5 频段功率特征
    """
    
    def __init__(self, fs: int = 250, channels: int = 14):
        self.fs = fs
        self.channels = channels
        # 频段定义 (Hz)
        self.bands = {
            'delta': (0.5, 4),
            'theta': (4, 8),
            'alpha': (8, 13),
            'beta': (13, 30),
            'gamma': (30, 45)
        }
    
    def extract(self, eeg_signal: np.ndarray) -> np.ndarray:
        """
        提取频域特征
        
        Args:
            eeg_signal: (C, T) 多通道 EEG 信号
            
        Returns:
            features: (C, 5) 各通道各频段功率
        """
        from scipy.signal import welch
        
        n_channels = eeg_signal.shape[0]
        features = np.zeros((n_channels, len(self.bands)))
        
        for ch in range(n_channels):
            freqs, psd = welch(eeg_signal[ch], fs=self.fs, nperseg=self.fs*2)
            
            for i, (band_name, (f_low, f_high)) in enumerate(self.bands.items()):
                mask = (freqs >= f_low) & (freqs <= f_high)
                features[ch, i] = np.sum(psd[mask])
        
        # 对数变换稳定方差
        features = np.log10(features + 1e-10)
        
        return features


class EyeTrackingFeatureExtractor:
    """
    眼动跟踪特征提取
    
    从原始眼动数据提取行为特征
    """
    
    def __init__(self, fps: int = 60):
        self.fps = fps
    
    def extract(self, gaze_data: np.ndarray) -> np.ndarray:
        """
        提取眼动特征
        
        Args:
            gaze_data: (T, 4) 时序眼动数据 [x, y, pupil, blink]
            
        Returns:
            features: (T, 6) 6维眼动特征
        """
        T = gaze_data.shape[0]
        features = np.zeros((T, 6))
        
        x, y, pupil, blink = gaze_data.T
        
        # 特征1: 注视持续时间（基于速度判定）
        velocity = np.sqrt(np.diff(x, prepend=x[0])**2 + 
                          np.diff(y, prepend=y[0])**2)
        fixation = (velocity < 0.01).astype(float)
        features[:, 0] = fixation
        
        # 特征2: 扫视频率（滑窗内扫视次数）
        window = int(self.fps * 2)  # 2秒窗口
        for i in range(T):
            start = max(0, i - window)
            features[i, 1] = np.sum(velocity[start:i+1] > 0.05) / (window / self.fps)
        
        # 特征3: 注视点分布熵
        for i in range(T):
            start = max(0, i - window)
            hist, _ = np.histogram2d(x[start:i+1], y[start:i+1], bins=10)
            p = hist.flatten() / (hist.sum() + 1e-10)
            features[i, 2] = -np.sum(p * np.log2(p + 1e-10))
        
        # 特征4: 瞳孔直径变化
        features[:, 3] = pupil
        
        # 特征5: 眨眼频率
        for i in range(T):
            start = max(0, i - window)
            features[i, 4] = np.sum(blink[start:i+1]) / (window / self.fps)
        
        # 特征6: PERCLOS
        for i in range(T):
            start = max(0, i - int(self.fps * 60))  # 60秒窗口
            features[i, 5] = np.mean(blink[start:i+1] > 0.5)
        
        return features


if __name__ == "__main__":
    # 模型初始化
    model = EEGEyeFusionModel(
        eeg_channels=14,
        eeg_bands=5,
        et_features=6,
        num_classes=3,
        hidden_dim=128
    )
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"融合模型参数量: {total_params:,}")
    
    # 模拟输入
    B = 8  # 批次
    eeg = torch.randn(B, 14, 5)    # 14通道×5频段
    et = torch.randn(B, 30, 6)     # 30时间步×6特征
    
    # 前向传播
    with torch.no_grad():
        logits = model(eeg, et)
    
    print(f"EEG 输入: {eeg.shape}")
    print(f"ET 输入: {et.shape}")
    print(f"输出: {logits.shape}")
    print(f"预测: {torch.argmax(logits, dim=1).tolist()}")
    
    # 特征提取器测试
    eeg_extractor = EEGFeatureExtractor(fs=250, channels=14)
    raw_eeg = np.random.randn(14, 2500)  # 10秒@250Hz
    eeg_feat = eeg_extractor.extract(raw_eeg)
    print(f"\nEEG 特征形状: {eeg_feat.shape}")
    print(f"频段功率示例 (通道0): δ={eeg_feat[0,0]:.2f}, "
          f"θ={eeg_feat[0,1]:.2f}, α={eeg_feat[0,2]:.2f}, "
          f"β={eeg_feat[0,3]:.2f}, γ={eeg_feat[0,4]:.2f}")
    
    et_extractor = EyeTrackingFeatureExtractor(fps=60)
    raw_et = np.random.randn(600, 4)  # 10秒@60Hz
    et_feat = et_extractor.extract(raw_et)
    print(f"ET 特征形状: {et_feat.shape}")
    
    print("\n=== 综述关键结论 ===")
    print("1. EEG+ET 融合在所有条件下一致提升分类性能")
    print("2. 跨被试准确率达 83%，跨性别达 87%")
    print("3. 注意力机制融合策略性能最优")
    print("4. 单眼动在某些场景优于单EEG（视觉分心/疲劳）")
    print("5. EEG在认知负荷/情绪检测方面更有优势")