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

class MultimodalFatigueDetector(nn.Module):
    """
    多模态被动疲劳检测器
    
    论文: Applied Sciences 2026, 16, 9049
    输入: 眼动特征 + ECG心电特征
    输出: 疲劳等级(清醒/轻度疲劳/严重疲劳)
    
    核心创新:
    1. 眼动特征: 注视模式+眨眼+瞳孔直径
    2. ECG特征: HRV心率变异性+频域分析
    3. 跨模态注意力: 自动调整两种模态权重
    """
    
    def __init__(self, eye_feat_dim=128, ecg_feat_dim=128,
                 hidden_dim=256, num_classes=3):
        super().__init__()
        
        # 1. 眼动特征编码器
        self.eye_encoder = EyeMovementEncoder(feat_dim=eye_feat_dim)
        
        # 2. ECG特征编码器
        self.ecg_encoder = ECGEncoder(feat_dim=ecg_feat_dim)
        
        # 3. 跨模态注意力融合
        self.cross_attn = CrossModalAttention(
            eye_dim=eye_feat_dim,
            ecg_dim=ecg_feat_dim,
            hidden_dim=hidden_dim
        )
        
        # 4. 疲劳分类头
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.GELU(),
            nn.Dropout(0.3),
            nn.Linear(hidden_dim // 2, num_classes)
        )
    
    def forward(self, eye_data, ecg_data):
        """
        Args:
            eye_data: dict with 'gaze', 'blink', 'pupil' tensors
            ecg_data: (B, T) ECG信号序列
        Returns:
            logits: (B, num_classes) 疲劳等级
            attn_weights: 跨模态注意力权重
        """
        # 编码
        eye_feat = self.eye_encoder(eye_data)  # (B, eye_feat_dim)
        ecg_feat = self.ecg_encoder(ecg_data)  # (B, ecg_feat_dim)
        
        # 跨模态融合
        fused, attn_weights = self.cross_attn(eye_feat, ecg_feat)
        
        # 分类
        logits = self.classifier(fused)
        
        return logits, attn_weights


class EyeMovementEncoder(nn.Module):
    """眼动特征编码器"""
    
    def __init__(self, feat_dim=128):
        super().__init__()
        
        # 注视方向序列编码
        self.gaze_encoder = nn.LSTM(
            input_size=2,  # (x, y)注视坐标
            hidden_size=64,
            num_layers=2,
            batch_first=True,
            bidirectional=True
        )
        
        # 眨眼特征
        self.blink_encoder = nn.Sequential(
            nn.Linear(5, 32),  # blink_rate, duration, interval, PERCLOS, avg_amplitude
            nn.GELU(),
            nn.Linear(32, 32)
        )
        
        # 瞳孔直径编码
        self.pupil_encoder = nn.Sequential(
            nn.Linear(3, 32),  # mean, std, trend
            nn.GELU(),
            nn.Linear(32, 32)
        )
        
        # 融合
        self.fusion = nn.Sequential(
            nn.Linear(128 + 32 + 32, feat_dim),  # LSTM bidirectional=128
            nn.GELU(),
            nn.Linear(feat_dim, feat_dim)
        )
    
    def forward(self, eye_data):
        gaze = eye_data['gaze']  # (B, T, 2)
        blink = eye_data['blink']  # (B, 5)
        pupil = eye_data['pupil']  # (B, 3)
        
        # 注视LSTM
        gaze_out, _ = self.gaze_encoder(gaze)
        gaze_feat = gaze_out[:, -1, :]  # (B, 128)
        
        # 眨眼
        blink_feat = self.blink_encoder(blink)  # (B, 32)
        
        # 瞳孔
        pupil_feat = self.pupil_encoder(pupil)  # (B, 32)
        
        # 融合
        combined = torch.cat([gaze_feat, blink_feat, pupil_feat], dim=-1)
        return self.fusion(combined)


class ECGEncoder(nn.Module):
    """ECG心电特征编码器"""
    
    def __init__(self, feat_dim=128, signal_len=500):
        super().__init__()
        
        # 1D卷积提取波形特征
        self.conv = nn.Sequential(
            nn.Conv1d(1, 32, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm1d(32), nn.ReLU(),
            nn.MaxPool1d(2),
            nn.Conv1d(32, 64, kernel_size=5, stride=2, padding=2),
            nn.BatchNorm1d(64), nn.ReLU(),
            nn.MaxPool1d(2),
            nn.Conv1d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm1d(128), nn.ReLU(),
            nn.AdaptiveAvgPool1d(1)
        )
        
        # HRV频域特征
        self.hrv_encoder = nn.Sequential(
            nn.Linear(6, 32),  # LF, HF, LF/HF, mean_HR, std_HR, RMSSD
            nn.GELU(),
            nn.Linear(32, 32)
        )
        
        # 融合
        self.fusion = nn.Sequential(
            nn.Linear(128 + 32, feat_dim),
            nn.GELU(),
            nn.Linear(feat_dim, feat_dim)
        )
    
    def forward(self, ecg_signal):
        """
        Args:
            ecg_signal: (B, T) ECG信号
        """
        # 波形特征
        x = ecg_signal.unsqueeze(1)  # (B, 1, T)
        wave_feat = self.conv(x).squeeze(-1)  # (B, 128)
        
        # HRV特征（假设已计算）
        hrv_feat = torch.randn(ecg_signal.shape[0], 6)  # 实际应从ECG提取
        hrv_encoded = self.hrv_encoder(hrv_feat)
        
        combined = torch.cat([wave_feat, hrv_encoded], dim=-1)
        return self.fusion(combined)


class CrossModalAttention(nn.Module):
    """跨模态注意力融合"""
    
    def __init__(self, eye_dim, ecg_dim, hidden_dim):
        super().__init__()
        
        # 投影到统一维度
        self.eye_proj = nn.Linear(eye_dim, hidden_dim)
        self.ecg_proj = nn.Linear(ecg_dim, hidden_dim)
        
        # 互注意力
        self.eye_to_ecg_attn = nn.MultiheadAttention(hidden_dim, num_heads=4, batch_first=True)
        self.ecg_to_eye_attn = nn.MultiheadAttention(hidden_dim, num_heads=4, batch_first=True)
        
        # 门控融合
        self.gate = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.Sigmoid()
        )
        
        self.output_proj = nn.Linear(hidden_dim * 2, hidden_dim)
    
    def forward(self, eye_feat, ecg_feat):
        # 投影
        eye = self.eye_proj(eye_feat).unsqueeze(1)  # (B, 1, hidden)
        ecg = self.ecg_proj(ecg_feat).unsqueeze(1)
        
        # 互注意力
        eye_attended, _ = self.eye_to_ecg_attn(eye, ecg, ecg)
        ecg_attended, _ = self.ecg_to_eye_attn(ecg, eye, eye)
        
        # 门控融合
        combined = torch.cat([eye_attended.squeeze(1), ecg_attended.squeeze(1)], dim=-1)
        gate_weight = self.gate(combined)
        fused = self.output_proj(combined * gate_weight)
        
        attn_weights = torch.stack([
            gate_weight.mean(dim=-1),
            (1 - gate_weight).mean(dim=-1)
        ], dim=-1)
        
        return fused, attn_weights


# 测试
if __name__ == "__main__":
    model = MultimodalFatigueDetector()
    
    # 模拟输入
    batch = 4
    eye_data = {
        'gaze': torch.randn(batch, 300, 2),  # 300帧注视
        'blink': torch.randn(batch, 5),
        'pupil': torch.randn(batch, 3)
    }
    ecg_signal = torch.randn(batch, 500)  # 5秒@100Hz
    
    logits, attn = model(eye_data, ecg_signal)
    print(f"Eye gaze: {eye_data['gaze'].shape}")
    print(f"ECG signal: {ecg_signal.shape}")
    print(f"Fatigue logits: {logits.shape}")
    print(f"Attention [eye, ecg]: {attn.shape}")
    print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")