from dataclasses import dataclass
from enum import Enum

class CognitiveState(Enum):
    """乘客认知状态（PCM 模型）"""
    NORMAL = "normal"           # 正常状态
    RISK_PREDICTION = "rp"     # 风险预测：提前感知潜在危险
    DANGER_IDENTIFICATION = "di"  # 危险识别：明确识别到危险

@dataclass
class PassengerCognitiveModel:
    """
    PCM: 乘客认知模型
    
    论文 Section 3.1: 区分两种认知过程
    
    Risk Prediction (RP):
    - 乘客在危险事件发生前就产生的预感
    - EEG 表现为前额叶 theta 频段能量上升
    - 时间窗口：事件前 2-5 秒
    
    Danger Identification (DI):
    - 乘客明确识别到危险事件
    - EEG 表现为 P300 等成分
    - 时间窗口：事件发生后 0.3-1 秒
    """
    
    sampling_rate: int = 250  # Hz
    num_channels: int = 32
    
    # RP 时间窗口
    rp_pre_event_sec: float = 5.0
    rp_post_event_sec: float = 0.0
    
    # DI 时间窗口
    di_pre_event_sec: float = 0.0
    di_post_event_sec: float = 1.0
    
    def get_rp_window(self, event_timestamp: float) -> tuple:
        """获取 RP 时间窗口"""
        start = event_timestamp - self.rp_pre_event_sec
        end = event_timestamp + self.rp_post_event_sec
        return start, end
    
    def get_di_window(self, event_timestamp: float) -> tuple:
        """获取 DI 时间窗口"""
        start = event_timestamp + self.di_pre_event_sec
        end = event_timestamp + self.di_post_event_sec
        return start, end
    
    def extract_frequency_bands(self, eeg_signal):
        """
        提取 5 个频段特征
        
        Returns:
            dict: 各频段能量
        """
        import numpy as np
        from scipy.signal import welch
        
        freqs, psd = welch(eeg_signal, fs=self.sampling_rate, nperseg=256)
        
        bands = {
            'delta': (1, 4),    # 慢波
            'theta': (4, 8),    # 认知负荷
            'alpha': (8, 13),   # 放松
            'beta': (13, 30),   # 活跃思维
            'gamma': (30, 45),  # 高级认知
        }
        
        band_powers = {}
        for name, (low, high) in bands.items():
            mask = (freqs >= low) & (freqs < high)
            band_powers[name] = np.mean(psd[mask]) if mask.any() else 0
        
        return band_powers


# 测试
if __name__ == "__main__":
    pcm = PassengerCognitiveModel()
    
    event_t = 10.0  # 事件时间戳
    rp_window = pcm.get_rp_window(event_t)
    di_window = pcm.get_di_window(event_t)
    
    print(f"Risk Prediction window: {rp_window}")
    print(f"Danger Identification window: {di_window}")
    
    # 模拟 EEG 信号
    import numpy as np
    eeg = np.random.randn(250 * 10)  # 10秒信号
    bands = pcm.extract_frequency_bands(eeg)
    print(f"Frequency bands: {bands}")

import torch
import torch.nn as nn

class Conv3DBlock(nn.Module):
    """3D 卷积块"""
    
    def __init__(self, in_ch, out_ch, kernel_size=3, stride=1, padding=1):
        super().__init__()
        self.conv = nn.Conv3d(in_ch, out_ch, kernel_size, stride, padding)
        self.bn = nn.BatchNorm3d(out_ch)
        self.relu = nn.ReLU()
        
    def forward(self, x):
        return self.relu(self.bn(self.conv(x)))


class RNNBlock(nn.Module):
    """双向 GRU 循环块"""
    
    def __init__(self, input_dim, hidden_dim, num_layers=2):
        super().__init__()
        self.bigru = nn.GRU(
            input_dim, hidden_dim, num_layers,
            batch_first=True, bidirectional=True
        )
        
    def forward(self, x):
        out, _ = self.bigru(x)
        return out


class PassengerEEGDecoder(nn.Module):
    """
    PEDS: 3D-CRNN 联合 EEG 解码模型
    
    论文 Section 3.3: 3D-CRNN
    
    架构:
    1. 3D CNN: 提取时空特征 (channels × time × freq)
    2. ResNet: 残差连接增强梯度流
    3. 双向 GRU: 捕获时序依赖
    4. 全连接: 分类输出
    
    输入: (B, 1, C, T, F) = (batch, 1, channels, time_steps, freq_bins)
    输出: (B, 2) = (batch, rp_prob, di_prob)
    """
    
    def __init__(self, num_channels=32, num_freq_bins=8, 
                 seq_len=125, hidden_dim=128, num_classes=2):
        super().__init__()
        
        # 3D CNN 特征提取
        self.conv1 = Conv3DBlock(1, 32, kernel_size=3, padding=1)
        self.conv2 = Conv3DBlock(32, 64, kernel_size=3, stride=2, padding=1)
        self.conv3 = Conv3DBlock(64, 128, kernel_size=3, stride=2, padding=1)
        self.pool = nn.AdaptiveAvgPool3d((None, 1, 1))
        
        # RNN 时序建模
        self.rnn = RNNBlock(128, hidden_dim, num_layers=2)
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim * 2, 64),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64, num_classes)
        )
        
    def forward(self, x):
        """
        Args:
            x: (B, 1, C, T, F) EEG 时空-频谱表示
            
        Returns:
            out: (B, 2) [RP_prob, DI_prob]
        """
        # 3D CNN
        x = self.conv1(x)   # (B, 32, C, T, F)
        x = self.conv2(x)   # (B, 64, C', T', F')
        x = self.conv3(x)   # (B, 128, C'', T'', F'')
        x = self.pool(x)    # (B, 128, T'', 1, 1)
        x = x.squeeze(-1).squeeze(-1)  # (B, 128, T'')
        
        # 转置为 (B, T'', 128)
        x = x.transpose(1, 2)
        
        # 双向 GRU
        x = self.rnn(x)    # (B, T'', hidden_dim*2)
        
        # 取最后一个时间步
        x = x[:, -1, :]    # (B, hidden_dim*2)
        
        # 分类
        out = self.classifier(x)  # (B, 2)
        
        return out


# 测试
if __name__ == "__main__":
    model = PassengerEEGDecoder(
        num_channels=32, num_freq_bins=8, seq_len=125
    )
    
    # 模拟输入: batch=4, 通道=32, 时间=125(0.5s@250Hz), 频率bins=8
    x = torch.randn(4, 1, 32, 125, 8)
    
    out = model(x)
    print(f"Input shape: {x.shape}")
    print(f"Output shape: {out.shape}")  # (4, 2)
    print(f"RP probs: {torch.softmax(out, dim=-1)[:, 0]}")
    print(f"DI probs: {torch.softmax(out, dim=-1)[:, 1]}")
    print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

import numpy as np
from typing import List, Tuple

class RiskAwareSequentialLabeling:
    """
    RSL: 风险感知顺序标注
    
    论文 Section 3.2
    
    核心思想:
    - 传统方法：事件前=正常(0)，事件后=危险(1)
    - RSL：在事件前一段时间就标注为"风险预测"(软标签)
    - 使用渐进式标签：0 → 0.3 → 0.7 → 1.0
    
    这提升了模型对"即将发生危险"的敏感性
    """
    
    def __init__(self, 
                 rp_window_sec: float = 5.0,
                 di_window_sec: float = 1.0,
                 sampling_rate: int = 250):
        self.rp_window = int(rp_window_sec * sampling_rate)
        self.di_window = int(di_window_sec * sampling_rate)
        self.sr = sampling_rate
        
    def label_sequence(self, 
                       total_length: int,
                       event_timestamps: List[float]) -> np.ndarray:
        """
        生成 RSL 标签序列
        
        Args:
            total_length: 总样本数
            event_timestamps: 事件时间戳列表（秒）
            
        Returns:
            labels: (N,) 0-1 之间的标签值
        """
        labels = np.zeros(total_length, dtype=np.float32)
        
        for t in event_timestamps:
            t_sample = int(t * self.sr)
            
            # DI 窗口：事件后，标签=1.0
            di_start = t_sample
            di_end = min(t_sample + self.di_window, total_length)
            labels[di_start:di_end] = 1.0
            
            # RP 窗口：事件前，渐进式标签
            rp_end = t_sample
            rp_start = max(0, t_sample - self.rp_window)
            
            if rp_start < rp_end:
                # 渐进式标签：从 0.1 渐增到 0.7
                rp_length = rp_end - rp_start
                rp_labels = np.linspace(0.1, 0.7, rp_length)
                labels[rp_start:rp_end] = np.maximum(
                    labels[rp_start:rp_end], 
                    rp_labels
                )
        
        return labels
    
    def generate_hard_labels(self, soft_labels: np.ndarray) -> np.ndarray:
        """将软标签转为硬标签用于评估"""
        hard = np.zeros_like(soft_labels, dtype=int)
        hard[soft_labels >= 0.5] = 1
        return hard


# 测试
if __name__ == "__main__":
    rsl = RiskAwareSequentialLabeling(
        rp_window_sec=5.0, di_window_sec=1.0, sampling_rate=250
    )
    
    # 60秒信号，2个事件：20s 和 40s
    total = 60 * 250
    events = [20.0, 40.0]
    
    labels = rsl.label_sequence(total, events)
    hard_labels = rsl.generate_hard_labels(labels)
    
    print(f"Total samples: {total}")
    print(f"RP samples (0.1-0.7): {((labels >= 0.1) & (labels < 0.8)).sum()}")
    print(f"DI samples (1.0): {(labels == 1.0).sum()}")
    print(f"Normal samples (0): {(labels == 0).sum()}")
    print(f"Hard positive samples: {hard_labels.sum()}")

def non_contact_eeg_feasibility():
    """
    非接触式 EEG 未来展望
    
    当前技术:
    - 座椅靠垫嵌入式电极: 信噪比低, 需要信号增强
    - 耳道电极: 嵌入耳机, 接触良好
    - 头枕电容耦合: 完全非接触, 但信号极弱
    
    预期时间线:
    2026-2028: 干电极头带（需用户佩戴）
    2028-2030: 座椅嵌入式湿电极（无需佩戴）
    2030+: 非接触电容耦合（完全无感）
    """
    timeline = {
        "2026-2028": {
            "技术": "干电极头带（Muse/Emotiv 级别）",
            "通道数": "2-5",
            "精度": "75-85%",
            "用户体验": "需佩戴头带",
            "量产可行性": "⚠️ 低（用户接受度问题）"
        },
        "2028-2030": {
            "技术": "座椅靠垫嵌入式电极",
            "通道数": "3-8",
            "精度": "70-80%",
            "用户体验": "无感",
            "量产可行性": "🟡 中（信号稳定性待验证）"
        },
        "2030+": {
            "技术": "非接触电容耦合 EEG",
            "通道数": "5-16",
            "精度": "65-75%",
            "用户体验": "完全无感",
            "量产可行性": "🔴 高（技术成熟后）"
        }
    }
    
    for period, info in timeline.items():
        print(f"\n{period}:")
        for k, v in info.items():
            print(f"  {k}: {v}")
    
    return timeline


if __name__ == "__main__":
    non_contact_eeg_feasibility()