"""
EEG 预处理与频谱特征提取

论文使用 Cognitive State Discrimination Dataset
包含多通道 EEG 采集，按认知状态标注
"""

import numpy as np
import torch
import torch.nn as nn
from scipy.signal import stft
from scipy.filters import butter, filtfilt

class EEGPreprocessor:
    """
    EEG 预处理管道
    
    步骤:
    1. 带通滤波 (0.5-45Hz)
    2. 工频陷波 (50/60Hz)
    3. 短时傅里叶变换 (STFT) → 频谱图
    4. 频带提取 (δ, θ, α, β, γ)
    """
    
    # 标准脑电频带定义
    BANDS = {
        'delta': (0.5, 4),    # δ: 深睡眠
        'theta': (4, 8),      # θ: 浅睡眠/冥想/记忆
        'alpha': (8, 13),     # α: 放松/清醒
        'beta': (13, 30),     # β: 活跃思考/专注
        'gamma': (30, 45),    # γ: 高级认知处理
    }
    
    def __init__(self, fs=250, n_channels=14):
        self.fs = fs
        self.n_channels = n_channels
        
    def preprocess(self, eeg_raw):
        """
        Args:
            eeg_raw: (n_channels, n_samples) 原始EEG信号
        Returns:
            spectrogram: (n_channels, n_freqs, n_times) 频谱图
            band_powers: (n_channels, 5) 五个频带能量
        """
        # 1. 带通滤波 0.5-45 Hz
        eeg_filtered = self._bandpass(eeg_raw, 0.5, 45)
        
        # 2. 工频陷波
        eeg_filtered = self._notch(eeg_filtered, 50)
        
        # 3. STFT
        freqs, times, Zxx = stft(
            eeg_filtered, 
            fs=self.fs,
            nperseg=128,
            noverlap=64
        )
        spectrogram = np.abs(Zxx)  # (n_channels, n_freqs, n_times)
        
        # 4. 频带能量提取
        band_powers = np.zeros((self.n_channels, 5))
        for i, (band, (low, high)) in enumerate(self.BANDS.items()):
            mask = (freqs >= low) & (freqs < high)
            band_powers[:, i] = np.mean(
            spectrogram[:, mask, :], axis=(1, 2)
            )
        
        # 归一化
        band_powers = band_powers / (band_powers.sum(axis=1, keepdims=True) + 1e-8)
        
        return spectrogram, band_powers
    
    def _bandpass(self, signal, low, high):
        b, a = butter(4, [low, high], btype='band', fs=self.fs)
        return filtfilt(b, a, signal, axis=-1)
    
    def _notch(self, signal, freq):
        b, a = butter(2, [freq-2, freq+2], btype='bandstop', fs=self.fs)
        return filtfilt(b, a, signal, axis=-1)


# ============================================
# 频谱域重建模块
# ============================================

class SpectralReconstruction(nn.Module):
    """
    频谱域重建模块 (SDRGNN 核心)
    
    将多通道 EEG 频谱通过图卷积重建为认知状态表征
    
    核心思想:
    - 不同脑区间存在功能连接
    - 认知分心会改变这些功能连接模式
    - 图神经网络可以学习这些拓扑关系
    - 频谱域重建恢复被噪声掩盖的认知信号
    """
    
    # 标准10-20系统电极位置 (14通道)
    ELECTRODE_POSITIONS = {
        'F3': (0.3, 0.3), 'F4': (0.7, 0.3),
        'C3': (0.3, 0.5), 'C4': (0.7, 0.5),
        'P3': (0.3, 0.7), 'P4': (0.7, 0.7),
        'O1': (0.35, 0.9), 'O2': (0.65, 0.9),
        'T7': (0.1, 0.5), 'T8': (0.9, 0.5),
        'Fz': (0.5, 0.25), 'Cz': (0.5, 0.5),
        'Pz': (0.5, 0.75), 'Oz': (0.5, 0.95),
    }
    
    def __init__(self, n_channels=14, n_freqs=65, embed_dim=128):
        super().__init__()
        
        # 频谱特征嵌入
        self.spec_embed = nn.Linear(n_freqs, embed_dim)
        
        # 预定义图邻接矩阵（基于电极物理距离）
        self.adj = self._build_adjacency()
        self.adj = nn.Parameter(self.adj, requires_grad=True)
        
        # 频谱域图卷积层
        self.spectral_gcn = nn.ModuleList([
            SpectralGCNLayer(embed_dim, embed_dim) 
            for _ in range(3)
        ])
        
        # 重建解码器
        self.decoder = nn.Sequential(
            nn.Linear(embed_dim, embed_dim * 2),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(embed_dim * 2, n_freqs),
        )
        
    def _build_adjacency(self):
        """基于电极距离构建邻接矩阵"""
        positions = list(self.ELECTRODE_POSITIONS.values())
        n = len(positions)
        adj = np.zeros((n, n))
        for i in range(n):
            for j in range(n):
                dist = np.sqrt(
                    (positions[i][0] - positions[j][0])**2 +
                    (positions[i][1] - positions[j][1])**2
                )
                adj[i, j] = np.exp(-dist * 3)  # 高斯核
        return torch.tensor(adj, dtype=torch.float32)
    
    def forward(self, spectrogram):
        """
        Args:
            spectrogram: (B, n_channels, n_freqs, n_times)
        Returns:
            reconstructed: (B, n_channels, n_freqs) 重建频谱
            features: (B, embed_dim) 认知状态表征
        """
        B, C, F, T = spectrogram.shape
        
        # 平均时间维度
        x = spectrogram.mean(dim=-1)  # (B, C, F)
        
        # 频谱嵌入
        x = self.spec_embed(x)  # (B, C, D)
        
        # 图卷积
        for gcn in self.spectral_gcn:
            x = gcn(x, self.adj)  # (B, C, D)
        
        # 全局池化
        features = x.mean(dim=1)  # (B, D)
        
        # 频谱重建
        reconstructed = self.decoder(x)  # (B, C, F)
        
        return reconstructed, features


class SpectralGCNLayer(nn.Module):
    """
    频谱域图卷积层
    
    基于拉普拉斯矩阵的特征分解做图卷积
    """
    def __init__(self, in_dim, out_dim):
        super().__init__()
        # 可学习权重
        self.W = nn.Linear(in_dim, out_dim)
        # 频谱滤波器
        self.spectral_filter = nn.Parameter(
            torch.randn(in_dim) * 0.01
        )
        self.bn = nn.BatchNorm1d(out_dim)
        
    def forward(self, x, adj):
        """
        Args:
            x: (B, C, D) 节点特征
            adj: (C, C) 邻接矩阵
        Returns:
            out: (B, C, D_out)
        """
        # 归一化拉普拉斯
        D = torch.diag(adj.sum(dim=1))
        L = D - adj  # 拉普拉斯矩阵
        # 归一化
        D_inv_sqrt = torch.diag(
            1.0 / torch.sqrt(adj.sum(dim=1) + 1e-8)
        )
        L_norm = D_inv_sqrt @ L @ D_inv_sqrt
        
        # 图卷积: X' = L_norm @ X @ W
        x_graph = torch.matmul(L_norm, x)  # (B, C, D)
        out = self.W(x_graph)  # (B, C, D_out)
        out = self.bn(out.transpose(1, 2)).transpose(1, 2)
        out = F.relu(out)
        return out


# ============================================
# ECG/GSR 处理模块
# ============================================

class ECGProcessor(nn.Module):
    """
    ECG 心电信号 → HRV 特征 → 认知负荷指标
    
    HRV (Heart Rate Variability) 与认知负荷密切相关:
    - 高认知负荷 → LF/HF 比值升高
    - 认知分心 → RMSSD 降低
    """
    def __init__(self, embed_dim=128):
        super().__init__()
        self.feature_extract = nn.Sequential(
            nn.Linear(8, 64),  # 8个HRV特征
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, embed_dim),
            nn.ReLU()
        )
        
    def forward(self, hrv_features):
        """
        Args:
            hrv_features: (B, 8) [SDNN, RMSSD, pNN50, 
                                 LF, HF, LF_HF_ratio, 
                                 mean_HR, std_HR]
        """
        return self.feature_extract(hrv_features)


class GSRProcessor(nn.Module):
    """
    GSR 皮肤电导反应 → 情绪唤醒 → 分心辅助判断
    
    GSR 与情绪唤醒度正相关:
    - 突然分心事件 → GSR 峰值
    - 持续高 GSR → 焦虑/压力
    """
    def __init__(self, embed_dim=128):
        super().__init__()
        self.feature_extract = nn.Sequential(
            nn.Linear(5, 32),  # 5个GSR特征
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(32, embed_dim),
            nn.ReLU()
        )
        
    def forward(self, gsr_features):
        """
        Args:
            gsr_features: (B, 5) [mean, std, peaks, 
                                  rise_time, recovery_time]
        """
        return self.feature_extract(gsr_features)

"""
多模态融合分类器

将 EEG 频谱重建特征 + ECG HRV 特征 + GSR 特征融合
进行三分类: Active / Resting / Focused
"""

class MultiModalFusion(nn.Module):
    """
    多模态融合 + 三分类
    
    融合策略: 交叉注意力 + 残差连接
    """
    def __init__(self, embed_dim=128, num_classes=3):
        super().__init__()
        
        # 模态间交叉注意力
        self.eeg_ecg_attn = CrossModalAttention(embed_dim)
        self.eeg_gsr_attn = CrossModalAttention(embed_dim)
        self.ecg_gsr_attn = CrossModalAttention(embed_dim)
        
        # 模态权重学习
        self.modal_weights = nn.Sequential(
            nn.Linear(embed_dim * 3, embed_dim),
            nn.ReLU(),
            nn.Linear(embed_dim, 3),
            nn.Softmax(dim=-1)
        )
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.Linear(embed_dim * 3, embed_dim * 2),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(embed_dim * 2, embed_dim),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(embed_dim, num_classes)
        )
        
    def forward(self, eeg_feat, ecg_feat, gsr_feat):
        """
        Args:
            eeg_feat: (B, D) EEG 频谱重建特征
            ecg_feat: (B, D) ECG HRV 特征
            gsr_feat: (B, D) GSR 特征
        Returns:
            logits: (B, 3) 三分类
            weights: (B, 3) 模态权重
        """
        # 交叉注意力
        eeg_enhanced = self.eeg_ecg_attn(eeg_feat, ecg_feat) + \
                       self.eeg_gsr_attn(eeg_feat, gsr_feat)
        ecg_enhanced = self.ecg_gsr_attn(ecg_feat, gsr_feat) + \
                       self.eeg_ecg_attn(ecg_feat, eeg_feat)
        gsr_enhanced = self.eeg_gsr_attn(gsr_feat, eeg_feat) + \
                       self.ecg_gsr_attn(gsr_feat, ecg_feat)
        
        # 模态权重
        combined = torch.cat([eeg_enhanced, ecg_enhanced, gsr_enhanced], dim=-1)
        weights = self.modal_weights(combined)
        
        # 加权融合
        weighted = (
            eeg_enhanced * weights[:, 0:1] +
            ecg_enhanced * weights[:, 1:2] +
            gsr_enhanced * weights[:, 2:3]
        )
        
        # 分类
        logits = self.classifier(combined)
        
        return logits, weights


class CrossModalAttention(nn.Module):
    """跨模态注意力"""
    def __init__(self, dim):
        super().__init__()
        self.query = nn.Linear(dim, dim)
        self.key = nn.Linear(dim, dim)
        self.value = nn.Linear(dim, dim)
        self.scale = dim ** -0.5
        
    def forward(self, query_feat, key_feat):
        q = self.query(query_feat).unsqueeze(1)  # (B, 1, D)
        k = self.key(key_feat).unsqueeze(1)      # (B, 1, D)
        v = self.value(key_feat).unsqueeze(1)    # (B, 1, D)
        
        attn = torch.softmax(
            torch.matmul(q, k.transpose(-1, -2)) * self.scale,
            dim=-1
        )
        out = torch.matmul(attn, v).squeeze(1)
        return out


# ============================================
# 完整模型
# ============================================

class SDRGNN(nn.Module):
    """
    MPS-DCDD-SDRGNN 完整模型
    
    频谱域重建图神经网络用于多模态认知分心检测
    """
    def __init__(self, n_eeg_channels=14, n_freqs=65, embed_dim=128):
        super().__init__()
        
        # EEG: 频谱域重建 + 图卷积
        self.spectral_recon = SpectralReconstruction(
            n_channels=n_eeg_channels,
            n_freqs=n_freqs,
            embed_dim=embed_dim
        )
        
        # ECG: HRV 特征提取
        self.ecg_processor = ECGProcessor(embed_dim=embed_dim)
        
        # GSR: 皮肤电导特征
        self.gsr_processor = GSRProcessor(embed_dim=embed_dim)
        
        # 多模态融合分类
        self.fusion = MultiModalFusion(
            embed_dim=embed_dim,
            num_classes=3
        )
        
    def forward(self, eeg_spectrogram, hrv_features, gsr_features):
        """
        Args:
            eeg_spectrogram: (B, 14, 65, T) EEG频谱
            hrv_features: (B, 8) 心率变异性特征
            gsr_features: (B, 5) 皮肤电导特征
        Returns:
            logits: (B, 3) 三分类
            reconstructed_spec: (B, 14, 65) 重建频谱
            modal_weights: (B, 3) 模态权重
        """
        # EEG 频谱域处理
        reconstructed_spec, eeg_feat = self.spectral_recon(eeg_spectrogram)
        
        # ECG/GSR 处理
        ecg_feat = self.ecg_processor(hrv_features)
        gsr_feat = self.gsr_processor(gsr_features)
        
        # 多模态融合
        logits, modal_weights = self.fusion(eeg_feat, ecg_feat, gsr_feat)
        
        return logits, reconstructed_spec, modal_weights


# ============================================
# 训练与评估
# ============================================

def train_sdrgnn():
    """训练入口"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    model = SDRGNN().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)
    
    # 分类损失 + 重建损失
    cls_criterion = nn.CrossEntropyLoss()
    recon_criterion = nn.MSELoss()
    
    # 训练循环 (伪代码)
    for epoch in range(50):
        model.train()
        for batch in train_loader:
            eeg_spec = batch['eeg'].to(device)  # (B, 14, 65, T)
            hrv = batch['hrv'].to(device)        # (B, 8)
            gsr = batch['gsr'].to(device)        # (B, 5)
            labels = batch['label'].to(device)   # (B,) 0/1/2
            
            optimizer.zero_grad()
            logits, recon_spec, weights = model(eeg_spec, hrv, gsr)
            
            # 组合损失
            cls_loss = cls_criterion(logits, labels)
            recon_loss = recon_criterion(recon_spec, eeg_spec.mean(dim=-1))
            total_loss = cls_loss + 0.1 * recon_loss
            
            total_loss.backward()
            optimizer.step()
        
        # 验证
        model.eval()
        all_preds, all_labels = [], []
        with torch.no_grad():
            for batch in val_loader:
                eeg_spec = batch['eeg'].to(device)
                hrv = batch['hrv'].to(device)
                gsr = batch['gsr'].to(device)
                labels = batch['label']
                
                logits, _, _ = model(eeg_spec, hrv, gsr)
                preds = logits.argmax(dim=-1).cpu()
                all_preds.extend(preds.numpy())
                all_labels.extend(labels.numpy())
        
        acc = accuracy_score(all_labels, all_preds)
        print(f"Epoch {epoch+1}: Val Acc = {acc:.4f}")

if __name__ == "__main__":
    train_sdrgnn()