import torch
import torch.nn as nn
import torch.nn.functional as F

class BidirectionalDelta(nn.Module):
    """
    双向Delta模块
    
    将一阶时序差分分解为正向（激活）和负向（抑制）两个分量，
    显式建模驾驶疲劳过程中神经活动的非对称动态。
    
    论文Section 3.1: Bidirectional Delta Module
    
    Args:
        channels: EEG通道数
        temp_size: 时间步数
    """
    
    def __init__(self, channels: int, temp_size: int):
        super().__init__()
        self.channels = channels
        self.temp_size = temp_size
        
        # 正向Delta变换（捕获神经激活）
        self.pos_conv = nn.Conv2d(channels, channels, kernel_size=(1, 3), 
                                   padding=(0, 1), groups=channels)
        # 负向Delta变换（捕获神经抑制）
        self.neg_conv = nn.Conv2d(channels, channels, kernel_size=(1, 3), 
                                   padding=(0, 1), groups=channels)
        
        # 通道注意力门控
        self.gate = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(channels * 2, channels),
            nn.Sigmoid()
        )
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: EEG信号, shape=(B, C, 1, T)
            
        Returns:
            双向Delta特征, shape=(B, C, 1, T)
        """
        batch_size = x.size(0)
        
        # 计算一阶差分: Δx(t) = x(t) - x(t-1)
        delta = x[:, :, :, 1:] - x[:, :, :, :-1]
        # 前后补零保持尺寸
        delta = F.pad(delta, (1, 0), mode='constant', value=0)
        
        # 分解为正向（激活）和负向（抑制）
        pos_delta = F.relu(delta)          # 正值 = 神经活动增强
        neg_delta = F.relu(-delta)         # 负值取正 = 神经活动抑制
        
        # 各自经过独立的深度卷积
        pos_feat = self.pos_conv(pos_delta)
        neg_feat = self.neg_conv(neg_delta)
        
        # 通道注意力门控融合
        gate_input = torch.cat([
            pos_feat.mean(dim=(2, 3)),
            neg_feat.mean(dim=(2, 3))
        ], dim=1)
        gate_weight = self.gate(gate_input).view(batch_size, self.channels, 1, 1)
        
        # 加权融合
        output = gate_weight * pos_feat + (1 - gate_weight) * neg_feat
        
        # 残差连接
        return output + delta


class GatedTemporalConv(nn.Module):
    """
    门控时序卷积模块
    
    使用深度可分离卷积捕获每个EEG通道的长期时序依赖，
    通过门控机制选择性传递信息。
    
    论文Section 3.2: Gated Temporal Convolution Module
    
    Args:
        channels: 通道数
        kernel_size: 卷积核大小
        dropout: dropout率
    """
    
    def __init__(self, channels: int, kernel_size: int = 3, dropout: float = 0.2):
        super().__init__()
        padding = (kernel_size - 1) // 2
        
        # 深度时序卷积（每通道独立）
        self.depthwise_conv = nn.Conv2d(
            channels, channels, 
            kernel_size=(1, kernel_size),
            padding=(0, padding),
            groups=channels
        )
        
        # 逐点卷积（通道间交互）
        self.pointwise_conv = nn.Conv2d(channels, channels, kernel_size=1)
        
        # 门控机制
        self.gate_conv = nn.Conv2d(
            channels, channels,
            kernel_size=(1, kernel_size),
            padding=(0, padding),
            groups=channels
        )
        
        self.bn1 = nn.BatchNorm2d(channels)
        self.bn2 = nn.BatchNorm2d(channels)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: shape=(B, C, 1, T)
            
        Returns:
            shape=(B, C, 1, T)
        """
        # 深度可分离卷积
        out = self.depthwise_conv(x)
        out = self.bn1(out)
        out = F.elu(out)
        out = self.pointwise_conv(out)
        out = self.bn2(out)
        
        # 门控
        gate = torch.sigmoid(self.gate_conv(x))
        out = out * gate
        
        # 残差
        out = self.dropout(out)
        return out + x


class DeltaGateNet(nn.Module):
    """
    DeltaGateNet: 双向时序动态建模框架
    
    论文: Bidirectional Temporal Dynamics Modeling for 
          EEG-based Driving Fatigue Recognition
    arXiv: 2602.14071
    
    架构: [Input] → [BiDelta] → [GatedTCN] × N → [GAP] → [FC] → [Output]
    
    Args:
        channels: EEG通道数 (SEED-VIG: 17, SADT: 30)
        temp_size: 时间步数
        num_blocks: GatedTCN块数
        num_classes: 分类数 (2: 疲劳/清醒)
    """
    
    def __init__(self, channels: int = 17, temp_size: int = 1200,
                 num_blocks: int = 3, num_classes: int = 2):
        super().__init__()
        
        # 双向Delta模块
        self.bi_delta = BidirectionalDelta(channels, temp_size)
        
        # 堆叠GatedTCN块
        self.tcn_blocks = nn.ModuleList([
            GatedTemporalConv(channels, kernel_size=3, dropout=0.2)
            for _ in range(num_blocks)
        ])
        
        # 分类头
        self.gap = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(channels, 64),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64, num_classes)
        )
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: EEG信号, shape=(B, C, 1, T)
            
        Returns:
            logits: shape=(B, num_classes)
        """
        # 双向Delta
        x = self.bi_delta(x)
        
        # GatedTCN堆叠
        for block in self.tcn_blocks:
            x = block(x)
        
        # 全局平均池化 + 分类
        x = self.gap(x)
        x = self.classifier(x)
        
        return x


# ============ 完整测试代码 ============
if __name__ == "__main__":
    import numpy as np
    from sklearn.metrics import accuracy_score, f1_score
    
    # 模拟SEED-VIG数据参数
    BATCH_SIZE = 32
    CHANNELS = 17       # SEED-VIG: 17通道
    TEMP_SIZE = 1200     # 5分钟 @ 4Hz
    NUM_CLASSES = 2
    
    # 初始化模型
    model = DeltaGateNet(
        channels=CHANNELS,
        temp_size=TEMP_SIZE,
        num_blocks=3,
        num_classes=NUM_CLASSES
    )
    
    # 统计参数量
    total_params = sum(p.numel() for p in model.parameters())
    print(f"DeltaGateNet参数量: {total_params:,}")
    
    # 模拟EEG数据（正常清醒状态：低振幅高频）
    np.random.seed(42)
    eeg_normal = torch.randn(BATCH_SIZE, CHANNELS, 1, TEMP_SIZE) * 0.5
    
    # 模拟EEG数据（疲劳状态：高振幅低频）
    t = np.linspace(0, 10, TEMP_SIZE)
    fatigue_pattern = np.sin(2 * np.pi * 2 * t) * 0.8  # 2Hz慢波
    eeg_fatigue = torch.tensor(
        fatigue_pattern + np.random.randn(BATCH_SIZE, CHANNELS, 1, TEMP_SIZE) * 0.3,
        dtype=torch.float32
    )
    
    # 前向传播
    print("\n--- 正常状态EEG ---")
    out_normal = model(eeg_normal)
    pred_normal = torch.argmax(out_normal, dim=1)
    print(f"输入shape: {eeg_normal.shape}")
    print(f"输出shape: {out_normal.shape}")
    print(f"预测: 清醒={torch.sum(pred_normal==0)}, 疲劳={torch.sum(pred_normal==1)}")
    
    print("\n--- 疲劳状态EEG ---")
    out_fatigue = model(eeg_fatigue)
    pred_fatigue = torch.argmax(out_fatigue, dim=1)
    print(f"输入shape: {eeg_fatigue.shape}")
    print(f"输出shape: {out_fatigue.shape}")
    print(f"预测: 清醒={torch.sum(pred_fatigue==0)}, 疲劳={torch.sum(pred_fatigue==1)}")
    
    # 训练流程示例
    print("\n--- 训练流程 ---")
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
    
    # 混合数据
    x_mix = torch.cat([eeg_normal, eeg_fatigue], dim=0)
    y_mix = torch.cat([
        torch.zeros(BATCH_SIZE, dtype=torch.long),   # 正常=0
        torch.ones(BATCH_SIZE, dtype=torch.long)     # 疲劳=1
    ])
    
    model.train()
    for epoch in range(5):
        optimizer.zero_grad()
        logits = model(x_mix)
        loss = criterion(logits, y_mix)
        loss.backward()
        optimizer.step()
        scheduler.step()
        
        acc = accuracy_score(y_mix.numpy(), torch.argmax(logits, dim=1).detach().numpy())
        print(f"Epoch {epoch+1}: Loss={loss.item():.4f}, Acc={acc:.2%}")
    
    # 推理速度测试
    model.eval()
    import time
    with torch.no_grad():
        start = time.time()
        for _ in range(100):
            _ = model(eeg_normal[:1])
        elapsed = (time.time() - start) / 100
        print(f"\n单样本推理延迟: {elapsed*1000:.1f}ms")
        print(f"理论帧率: {1/elapsed:.0f} FPS")

def multimodal_fatigue_fusion(visual_score: float, eeg_score: float,
                                eeg_confidence: float = 0.8) -> float:
    """
    多模态疲劳融合：视觉PERCLOS + EEG DeltaGateNet
    
    Args:
        visual_score: PERCLOS百分比 (0-100)
        eeg_score: EEG模型疲劳概率 (0-1)
        eeg_confidence: EEG信号质量置信度 (0-1)
        
    Returns:
        fused_score: 融合后疲劳分数 (0-100)
    """
    # 视觉权重随光照/遮挡自适应
    visual_weight = 1.0 - eeg_confidence * 0.3
    
    # EEG权重随信号质量调整
    eeg_weight = eeg_confidence * 0.3
    
    # 归一化
    total = visual_weight + eeg_weight
    visual_weight /= total
    eeg_weight /= total
    
    fused = (visual_score * visual_weight + 
             eeg_score * 100 * eeg_weight)
    
    return fused

# 示例
print(multimodal_fatigue_fusion(
    visual_score=25.0,   # PERCLOS 25%
    eeg_score=0.78,       # EEG模型判断78%概率疲劳
    eeg_confidence=0.85   # EEG信号质量良好
))
# 输出: ~43.8% → 触发二级疲劳警告