# 示意代码，不可直接运行（非论文官方实现；论文实际为 CNN+BiLSTM+时间注意力编码，非 Transformer）
import torch
import torch.nn as nn

class EyeTrackingStream(nn.Module):
    """
    眼动追踪流：Transformer编码器
    
    提取时序眼动特征：
    - 凝视位置序列
    - 眨眼频率/时长
    - 扫视幅度/方向
    - 瞳孔直径变化
    """
    
    def __init__(self, d_model: int = 128, n_heads: int = 4,
                 n_layers: int = 3, input_dim: int = 8):
        super().__init__()
        # 输入投影：8维眼动特征→d_model
        self.input_proj = nn.Linear(input_dim, d_model)
        
        # 位置编码
        self.pos_enc = PositionalEncoding(d_model)
        
        # Transformer编码器
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=n_heads,
            dim_feedforward=d_model*4,
            dropout=0.1, batch_first=True,
            activation='gelu'
        )
        self.encoder = nn.TransformerEncoder(
            encoder_layer, num_layers=n_layers
        )
        
        # 时序池化
        self.temporal_pool = nn.AdaptiveAvgPool1d(1)
        
    def forward(self, eye_data: torch.Tensor) -> torch.Tensor:
        """
        Args:
            eye_data: [B, T, 8] 眼动特征序列
            特征维度: [gaze_x, gaze_y, pupil_diameter,
                       blink_flag, saccade_amp, 
                       saccade_vel, fixation_dur, 
                       gaze_dispersion]
            
        Returns:
            features: [B, d_model] 时序特征
        """
        x = self.input_proj(eye_data)  # [B, T, d_model]
        x = self.pos_enc(x)
        x = self.encoder(x)
        # 时序平均池化
        x = x.transpose(1, 2)  # [B, d_model, T]
        x = self.temporal_pool(x).squeeze(-1)  # [B, d_model]
        return x


class PositionalEncoding(nn.Module):
    """正弦位置编码"""
    def __init__(self, d_model: int, max_len: int = 3600):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * 
            (-torch.log(torch.tensor(10000.0)) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe.unsqueeze(0))
        
    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

# 示意代码，不可直接运行（非论文官方实现；论文的ECG流输入为ECG-derived特征而非原始波形，
# 专家特征仅用RMSSD）
class ECGStream(nn.Module):
    """
    ECG流：CNN编码器 + HRV专家特征
    
    双路径：
    1. CNN从原始ECG波形提取深度特征
    2. 专家特征：从ECG计算的HRV指标
    """
    
    def __init__(self, signal_length: int = 1800,  # 60s@30Hz
                 d_model: int = 128):
        super().__init__()
        # CNN编码器
        self.cnn = nn.Sequential(
            nn.Conv1d(1, 32, 7, stride=2, padding=3),
            nn.BatchNorm1d(32),
            nn.ReLU(),
            
            nn.Conv1d(32, 64, 5, stride=2, padding=2),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            
            nn.Conv1d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            
            nn.AdaptiveAvgPool1d(1),
            nn.Flatten(),
            nn.Linear(128, d_model)
        )
        
        # 专家特征维度
        self.expert_dim = 6
        # 专家特征投影
        self.expert_proj = nn.Sequential(
            nn.Linear(self.expert_dim, 32),
            nn.ReLU(),
            nn.Linear(32, d_model)
        )
        
    def extract_hrv_features(
        self, ecg: torch.Tensor, fs: int = 30
    ) -> torch.Tensor:
        """
        提取HRV专家特征
        
        Args:
            ecg: [B, L] ECG信号
            
        Returns:
            hrv: [B, 6] HRV特征
        """
        B = ecg.shape[0]
        device = ecg.device
        
        # R峰检测（简化版：使用阈值+找峰）
        r_peaks = self._detect_r_peaks(ecg, fs)
        
        features = []
        for b in range(B):
            peaks = r_peaks[b]
            n_peaks = len(peaks)
            
            if n_peaks < 3:
                features.append(torch.zeros(self.expert_dim, device=device))
                continue
            
            # RR间期
            rr = torch.diff(peaks.float()) / fs  # 秒
            
            # 1. 平均心率
            mean_hr = 60.0 / rr.mean() if rr.mean() > 0 else 0
            
            # 2. RMSSD
            rmssd = torch.sqrt(torch.mean(rr[1:] - rr[:-1] ** 2))
            
            # 3. SDNN
            sdnn = rr.std()
            
            # 4. LF功率 (0.04-0.15 Hz)
            # 5. HF功率 (0.15-0.4 Hz)
            # 6. LF/HF比值
            if len(rr) > 10:
                rr_interp = torch.nn.functional.interpolate(
                    rr.view(1, -1), size=256, mode='linear'
                ).squeeze(0)
                freqs = torch.fft.rfftfreq(256, d=1/fs)
                power = torch.abs(torch.fft.rfft(rr_interp)) ** 2
                
                lf_mask = (freqs >= 0.04) & (freqs < 0.15)
                hf_mask = (freqs >= 0.15) & (freqs <= 0.4)
                
                lf_power = power[lf_mask].sum()
                hf_power = power[hf_mask].sum()
                lf_hf = (lf_power / (hf_power + 1e-8)).item()
            else:
                lf_power = hf_power = lf_hf = 0
            
            features.append(torch.tensor([
                mean_hr, rmssd.item(), sdnn.item(),
                lf_power.item(), hf_power.item(), lf_hf
            ], device=device))
        
        return torch.stack(features)
    
    def _detect_r_peaks(self, ecg, fs):
        """简化版R峰检测"""
        # 实际应使用Pan-Tompkins算法
        peaks = []
        for b in range(ecg.shape[0]):
            signal = ecg[b].cpu().numpy()
            threshold = np.percentile(signal[signal > 0], 95)
            above = signal > threshold
            # 找连续段的中点
            from scipy.signal import find_peaks
            pks, _ = find_peaks(signal, height=threshold, 
                              distance=fs*0.3)
            peaks.append(torch.tensor(pks))
        return peaks
    
    def forward(self, ecg: torch.Tensor) -> torch.Tensor:
        """
        Args:
            ecg: [B, L] ECG信号
            
        Returns:
            features: [B, d_model] ECG特征
        """
        # CNN深度特征
        ecg_input = ecg.unsqueeze(1)  # [B, 1, L]
        cnn_feat = self.cnn(ecg_input)
        
        # HRV专家特征
        hrv = self.extract_hrv_features(ecg)
        expert_feat = self.expert_proj(hrv)
        
        # 融合
        return cnn_feat + expert_feat


class DualStreamFatigueDetector(nn.Module):
    """
    双流疲劳检测模型
    
    融合眼动+ECG进行疲劳分类
    """
    
    def __init__(self, n_classes: int = 3):
        super().__init__()
        self.eye_stream = EyeTrackingStream(d_model=128)
        self.ecg_stream = ECGStream(d_model=128)
        
        # 后期融合
        self.fusion = nn.Sequential(
            nn.Linear(128 + 128, 64),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64, n_classes)
        )
        
    def forward(self, eye_data: torch.Tensor, 
                ecg: torch.Tensor) -> torch.Tensor:
        eye_feat = self.eye_stream(eye_data)
        ecg_feat = self.ecg_stream(ecg)
        fused = torch.cat([eye_feat, ecg_feat], dim=1)
        return self.fusion(fused)


if __name__ == "__main__":
    model = DualStreamFatigueDetector(n_classes=3)
    # 5分钟@30Hz眼动 = 9000帧, 取30秒窗口
    eye_input = torch.randn(4, 900, 8)  # [B, T, 8]
    ecg_input = torch.randn(4, 1800)    # [B, 60s@30Hz]
    
    output = model(eye_input, ecg_input)
    print(f"眼动输入: {eye_input.shape}")
    print(f"ECG输入: {ecg_input.shape}")
    print(f"输出: {output.shape}")
    print(f"参数量: {sum(p.numel() for p in model.parameters()):,}")