import numpy as np
from scipy import signal as sig

class EarEEGDistractionDetector:
    """
    单通道耳道EEG驾驶员分心检测
    
    论文核心: bioRxiv 2026.01.24.701469
    关键指标: frontal theta升高 + alpha降低 = 认知分心
    
    优势:
    - 单通道（耳道放置），无需全头皮电极
    - 非侵入式，可集成到耳机/头枕
    - 实时计算，延迟<100ms
    """
    
    def __init__(self, fs=250):
        self.fs = fs  # 采样率
        # 频段定义
        self.bands = {
            'theta': (4, 8),    # 认知负荷
            'alpha': (8, 13),   # 放松/注意力
            'beta': (13, 30),   # 警觉
        }
        self.window_size = int(4 * fs)  # 4秒窗口
        self.threshold = 1.5  # theta/alpha比值阈值
    
    def compute_psd(self, eeg_signal):
        """
        计算功率谱密度
        
        Args:
            eeg_signal: shape=(N,), 单通道EEG
        
        Returns:
            freqs, psd: 频率和功率
        """
        freqs, psd = sig.welch(eeg_signal, self.fs, nperseg=self.window_size)
        return freqs, psd
    
    def band_power(self, freqs, psd, band):
        """计算指定频段的平均功率"""
        f_low, f_high = band
        mask = (freqs >= f_low) & (freqs <= f_high)
        return np.mean(psd[mask]) if np.any(mask) else 0
    
    def compute_theta_alpha_ratio(self, eeg_signal):
        """
        计算 theta/alpha 比值
        
        这是分心检测的核心指标：
        - 比值升高 → 认知负荷增加 → 可能分心
        - 比值降低 → 放松状态 → 正常驾驶
        
        Args:
            eeg_signal: 4秒EEG窗口, shape=(1000,) @250Hz
        
        Returns:
            ratio: theta/alpha功率比值
            is_distracted: 是否分心
        """
        freqs, psd = self.compute_psd(eeg_signal)
        
        theta_power = self.band_power(freqs, psd, self.bands['theta'])
        alpha_power = self.band_power(freqs, psd, self.bands['alpha'])
        
        # 避免除零
        ratio = theta_power / (alpha_power + 1e-10)
        
        is_distracted = ratio > self.threshold
        
        return ratio, is_distracted
    
    def sliding_window_detect(self, eeg_stream, step_sec=1):
        """
        滑动窗口实时检测
        
        Args:
            eeg_stream: 连续EEG信号
            step_sec: 步长（秒）
        
        Returns:
            results: [(time_sec, ratio, is_distracted), ...]
        """
        step = int(step_sec * self.fs)
        results = []
        
        for i in range(0, len(eeg_stream) - self.window_size, step):
            window = eeg_stream[i:i + self.window_size]
            ratio, distracted = self.compute_theta_alpha_ratio(window)
            t = i / self.fs
            results.append((t, ratio, distracted))
        
        return results


# 测试代码
if __name__ == "__main__":
    detector = EarEEGDistractionDetector(fs=250)
    
    # 模拟正常驾驶EEG（alpha主导）
    np.random.seed(42)
    t = np.linspace(0, 10, 2500)
    normal_eeg = (
        5 * np.sin(2 * np.pi * 10 * t) +   # alpha 10Hz
        1 * np.sin(2 * np.pi * 6 * t) +    # theta 6Hz
        0.5 * np.random.randn(2500)         # 噪声
    )
    
    # 模拟分心驾驶EEG（theta主导）
    distracted_eeg = (
        1 * np.sin(2 * np.pi * 10 * t) +   # alpha降低
        5 * np.sin(2 * np.pi * 6 * t) +     # theta升高
        0.5 * np.random.randn(2500)
    )
    
    ratio_normal, _ = detector.compute_theta_alpha_ratio(normal_eeg[:1000])
    ratio_distract, _ = detector.compute_theta_alpha_ratio(distracted_eeg[:1000])
    
    print(f"正常驾驶 theta/alpha: {ratio_normal:.3f} (阈值 {detector.threshold})")
    print(f"分心驾驶 theta/alpha: {ratio_distract:.3f} (阈值 {detector.threshold})")

class EarEEGHardware:
    """
    耳道EEG传感器硬件方案
    
    可集成到:
    - 车载头枕（座椅靠背两侧）
    - 智能耳机（驾驶辅助耳机）
    - AR眼镜（如Meta Aria Gen 2）
    """
    
    configs = {
        "in_ear": {
            "placement": "耳道内",
            "electrodes": 1,  # 单通道
            "reference": "耳垂",
            "snr": "中等",
            "comfort": "高（类似入耳式耳机）",
            "signal_quality": "可检测theta/alpha"
        },
        "around_ear": {
            "placement": "耳廓周围",
            "electrodes": 2,  # 差分
            "reference": "对侧耳",
            "snr": "较高",
            "comfort": "高（类似后挂式耳机）",
            "signal_quality": "更好的频段分离"
        },
        "headrest_embedded": {
            "placement": "头枕接触面",
            "electrodes": 2,
            "reference": "头枕金属框架",
            "snr": "中等（接触阻抗较高）",
            "comfort": "无需佩戴",
            "signal_quality": "需要信号增强"
        }
    }

class VAEOCSVM酒精检测:
    """
    半监督酒驾检测：VAE特征提取 + OCSVM异常检测
    
    优势：仅需正常驾驶数据训练，酒驾作为异常检测
    """
    
    def __init__(self, input_dim=20, latent_dim=8):
        # VAE编码器
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, latent_dim * 2)  # mean + log_var
        )
        # OCSVM
        self.ocsvm = OneClassSVM(kernel='rbf', nu=0.05)
    
    def detect(self, driving_features):
        """检测酒驾（异常）"""
        # VAE提取特征
        encoded = self.encode(driving_features)
        # OCSVM判断异常
        return self.ocsvm.predict(encoded) == -1  # -1=异常