"""
入耳式EEG信号采集与处理
适用于座舱内认知负荷监测
"""

import numpy as np
from scipy.signal import butter, filtfilt, welch

class InEarEEGProcessor:
    """
    入耳式EEG信号处理器
    
    硬件：入耳式电极（2-4通道）
    采样率：256-512 Hz
    应用：驾驶员认知负荷/疲劳连续监测
    """
    
    def __init__(self, config: dict):
        self.fs = config.get('sample_rate', 256)  # 采样率 Hz
        self.n_channels = config.get('n_channels', 2)  # 耳道电极数
        self.ref_channel = config.get('ref', 'linked_ears')  # 参考电极
        
        # 滤波器参数
        self.band_low = config.get('band_low', 0.5)   # 带通低频
        self.band_high = config.get('band_high', 45.0) # 带通高频
        self.notch_freq = config.get('notch', 50.0)   # 工频陷波
        
    def preprocess(self, raw_signal: np.ndarray) -> np.ndarray:
        """
        预处理：滤波 → 去伪迹 → 参考
        
        Args:
            raw_signal: 原始EEG信号, shape=(n_samples, n_channels)
        
        Returns:
            clean_signal: 清理后信号
        """
        # 带通滤波 (0.5-45 Hz)
        b, a = butter(4, [self.band_low/(self.fs/2), 
                          self.band_high/(self.fs/2)], 
                      btype='band')
        filtered = filtfilt(b, a, raw_signal, axis=0)
        
        # 工频陷波
        if self.notch_freq > 0:
            Q = 30.0
            w0 = self.notch_freq / (self.fs / 2)
            b_n, a_n = butter(2, w0, btype='notch')
            filtered = filtfilt(b_n, a_n, filtered, axis=0)
        
        # 去除眼电伪迹（简单阈值法）
        threshold = np.mean(np.abs(filtered)) + 5 * np.std(filtered)
        mask = np.abs(filtered) < threshold
        clean = filtered * mask
        
        return clean
    
    def extract_features(self, signal: np.ndarray) -> dict:
        """
        提取认知负荷相关特征
        
        频段定义：
        - Delta (δ): 0.5-4 Hz → 深睡眠
        - Theta (θ): 4-8 Hz → 疲劳/困倦
        - Alpha (α): 8-13 Hz → 放松/闭眼
        - Beta (β): 13-30 Hz → 活跃/专注
        - Gamma (γ): 30-45 Hz → 高级认知
        
        关键指标：
        - Theta/Alpha比：疲劳程度
        - Beta/Theta比：警觉水平
        - Alpha不对称：情感效价
        """
        # 功率谱密度
        freqs, psd = welch(signal, fs=self.fs, nperseg=self.fs*2)
        
        # 频段功率
        bands = {
            'delta': (0.5, 4),
            'theta': (4, 8),
            'alpha': (8, 13),
            'beta': (13, 30),
            'gamma': (30, 45)
        }
        
        band_power = {}
        for name, (low, high) in bands.items():
            mask = (freqs >= low) & (freqs < high)
            band_power[name] = np.mean(psd[mask]) if np.any(mask) else 0
        
        # 认知负荷指数
        theta_alpha_ratio = (band_power['theta'] + 1e-10) / (band_power['alpha'] + 1e-10)
        beta_theta_ratio = (band_power['beta'] + 1e-10) / (band_power['theta'] + 1e-10)
        
        # 归一化负荷指数 [0, 1]
        # 高Theta/Alpha = 高疲劳
        # 低Beta/Theta = 低警觉
        fatigue_index = np.clip(theta_alpha_ratio / 2.0, 0, 1)
        alertness_index = np.clip(beta_theta_ratio / 3.0, 0, 1)
        
        return {
            'band_power': band_power,
            'theta_alpha_ratio': theta_alpha_ratio,
            'beta_theta_ratio': beta_theta_ratio,
            'fatigue_index': fatigue_index,
            'alertness_index': alertness_index,
            'cognitive_load': 1.0 - alertness_index  # 负荷 = 1 - 警觉
        }
    
    def detect_microsleep(self, signal: np.ndarray) -> bool:
        """
        微睡眠检测：Alpha功率骤增 + Theta功率上升
        
        微睡眠特征：
        - 持续1-15秒
        - Alpha功率突然上升（闭眼效应）
        - Theta功率持续上升
        - 行为表现：短暂失神
        """
        features = self.extract_features(signal)
        
        alpha_power = features['band_power']['alpha']
        theta_power = features['band_power']['theta']
        
        # 微睡眠判据
        alpha_surge = alpha_power > np.mean(psd) * 2  # Alpha骤增
        theta_high = theta_power > np.mean(psd) * 1.5  # Theta偏高
        
        return alpha_surge and theta_high


# ============ 座舱内应用模拟 ============
if __name__ == "__main__":
    np.random.seed(42)
    
    # 模拟2分钟EEG数据（2通道，256Hz）
    n_samples = 120 * 256
    # 正常驾驶状态
    eeg_normal = np.random.normal(0, 10, (n_samples, 2))
    
    # 疲劳状态（Theta增加，Beta减少）
    t = np.arange(n_samples) / 256
    theta_wave = 5 * np.sin(2 * np.pi * 6 * t)  # 6Hz Theta
    eeg_fatigue = eeg_normal + theta_wave.reshape(-1, 1)
    
    processor = InEarEEGProcessor({'sample_rate': 256, 'n_channels': 2})
    
    # 正常状态分析
    clean_normal = processor.preprocess(eeg_normal)
    feat_normal = processor.extract_features(clean_normal)
    
    # 疲劳状态分析
    clean_fatigue = processor.preprocess(eeg_fatigue)
    feat_fatigue = processor.extract_features(clean_fatigue)
    
    print("=" * 50)
    print("入耳式EEG座舱内认知监测")
    print("=" * 50)
    print("\n正常驾驶状态:")
    print(f"  Theta/Alpha比: {feat_normal['theta_alpha_ratio']:.3f}")
    print(f"  Beta/Theta比: {feat_normal['beta_theta_ratio']:.3f}")
    print(f"  疲劳指数: {feat_normal['fatigue_index']:.3f}")
    print(f"  警觉指数: {feat_normal['alertness_index']:.3f}")
    
    print("\n疲劳驾驶状态:")
    print(f"  Theta/Alpha比: {feat_fatigue['theta_alpha_ratio']:.3f}")
    print(f"  Beta/Theta比: {feat_fatigue['beta_theta_ratio']:.3f}")
    print(f"  疲劳指数: {feat_fatigue['fatigue_index']:.3f}")
    print(f"  警觉指数: {feat_fatigue['alertness_index']:.3f}")
    
    print("\n" + "=" * 50)
    print("消费级入耳式EEG厂商对比:")
    
    vendors = [
        ("Naox Technologies", "法国", "FDA 510(k)", "$299", "医疗+消费"),
        ("Neurable", "美国", "CE认证", "$399", "疲劳管理"),
        ("NextSense", "美国", "FDA临床", "$499", "睡眠监测"),
        ("Idia", "丹麦", "CE认证", "$249", "冥想辅助"),
        ("Neurosteer", "以色列", "FDA临床", "$350", "麻醉监护"),
    ]
    
    print(f"{'厂商':<20} {'国家':<6} {'认证':<12} {'价格':<8} {'定位'}")
    print("-" * 65)
    for v in vendors:
        print(f"{v[0]:<20} {v[1]:<6} {v[2]:<12} {v[3]:<8} {v[4]}")

"""
车载入耳式EEG数据接入方案
"""

# BLE GATT Service定义
EEG_SERVICE_UUID = "0000fe3d-0000-1000-8000-00805f9b34fb"
EEG_DATA_CHAR_UUID = "0000fe3e-0000-1000-8000-00805f9b34fb"

# 数据协议
eeg_packet = {
    "header": 0xAA55,
    "timestamp_ms": 0,          # 毫秒时间戳
    "channels": 2,              # 通道数
    "sample_rate": 256,         # 采样率
    "resolution": 16,           # ADC分辨率
    "data": [0] * 512,          # 1秒数据（2ch × 256samples）
    "battery_pct": 85,          # 电量
    "impedance_ok": True,       # 电极阻抗
}

# 集成方案
integration_plan = {
    "phase1_poc": {
        "duration": "2周",
        "task": "BLE连接 + 原始数据接收 + 实时频谱显示",
        "hardware": "开发板 + Neurable耳机",
        "output": "数据质量评估报告"
    },
    "phase2_analysis": {
        "duration": "4周",
        "task": "疲劳指数计算 + PERCLOS对照验证",
        "hardware": "测试车辆 + Neurable + IR摄像头",
        "output": "EEG vs PERCLOS相关性分析"
    },
    "phase3_integration": {
        "duration": "8周",
        "task": "DMS融合决策 + 警告策略",
        "hardware": "Snapdragon 8295",
        "output": "多模态DMS原型"
    }
}

for phase, details in integration_plan.items():
    print(f"\n{phase}:")
    for k, v in details.items():
        print(f"  {k}: {v}")