import numpy as np
import mne
from scipy import signal

class MindWanderingDetector:
    """基于 EEG 的 Mind Wandering 检测器
    
    参考：Baldwin et al., 2017
    """
    
    def __init__(self, sfreq=250):
        self.sfreq = sfreq
        # Alpha 波段：8-14 Hz
        self.alpha_band = (8, 14)
        # 走神阈值：Alpha 功率 > 1.5 倍基线
        self.threshold = 1.5
        
    def extract_alpha_power(self, eeg_data):
        """提取 Alpha 波段功率
        
        Args:
            eeg_data: EEG 数据, shape=(n_channels, n_samples)
            
        Returns:
            alpha_power: Alpha 功率, shape=(n_channels,)
        """
        # 带通滤波 8-14 Hz
        sos = signal.butter(4, self.alpha_band, btype='band', 
                           fs=self.sfreq, output='sos')
        filtered = signal.sosfiltfilt(sos, eeg_data, axis=1)
        
        # 计算功率谱密度
        f, psd = signal.welch(filtered, fs=self.sfreq, nperseg=512)
        
        # 提取 Alpha 波段
        alpha_idx = np.where((f >= self.alpha_band[0]) & 
                             (f <= self.alpha_band[1]))[0]
        alpha_power = np.mean(psd[:, alpha_idx], axis=1)
        
        return alpha_power
    
    def detect_mind_wandering(self, eeg_window, baseline):
        """检测 Mind Wandering 状态
        
        Args:
            eeg_window: 当前 EEG 窗口, shape=(n_channels, window_samples)
            baseline: 基线 Alpha 功率, shape=(n_channels,)
            
        Returns:
            is_mind_wandering: bool, 是否处于走神状态
            confidence: float, 置信度 0-1
        """
        alpha_power = self.extract_alpha_power(eeg_window)
        
        # 相对功率比
        power_ratio = alpha_power / baseline
        
        # 顶叶和枕叶通道（Pz, O1, O2）
        parietal_occipital_idx = [29, 35, 36]  # 示例通道索引
        avg_ratio = np.mean(power_ratio[parietal_occipital_idx])
        
        # 判断
        is_mind_wandering = avg_ratio > self.threshold
        confidence = min(avg_ratio / self.threshold, 1.0)
        
        return is_mind_wandering, confidence


# 实际测试
if __name__ == "__main__":
    # 模拟数据
    np.random.seed(42)
    n_channels = 64
    sfreq = 250
    window_sec = 5
    
    # 生成模拟 EEG（专注状态）
    baseline_eeg = np.random.randn(n_channels, sfreq * 60) * 10  # μV
    
    # 生成模拟 EEG（走神状态，Alpha 功率增加）
    mind_wandering_eeg = baseline_eeg.copy()
    # 在 Alpha 波段增加能量
    t = np.arange(sfreq * window_sec) / sfreq
    alpha_wave = 15 * np.sin(2 * np.pi * 10 * t)  # 10 Hz
    mind_wandering_eeg[29, -sfreq*window_sec:] += alpha_wave
    
    # 检测
    detector = MindWanderingDetector(sfreq=sfreq)
    baseline = detector.extract_alpha_power(baseline_eeg[:, :sfreq*60])
    
    is_mw, conf = detector.detect_mind_wandering(
        mind_wandering_eeg[:, -sfreq*window_sec:], baseline
    )
    
    print(f"走神状态: {is_mw}")
    print(f"置信度: {conf:.2f}")

def detect_mind_wandering_behavioral(speed, lateral_position, steering_angle, 
                                     baseline_speed, baseline_sdlp, baseline_srr):
    """基于驾驶行为的 Mind Wandering 检测
    
    参考：He et al., 2011
    
    特征：
    - 速度波动增加
    - 车道偏差减少（过度平滑）
    - 方向盘反转率降低
    
    Returns:
        is_mw: bool, 是否走神
        score: float, 综合评分
    """
    # 计算当前特征
    speed_var = np.std(speed)
    sdlp = np.std(lateral_position)
    srr = len(np.where(np.diff(np.sign(steering_angle)))[0])
    
    # 相对变化
    speed_var_ratio = speed_var / np.std(baseline_speed)
    sdlp_ratio = sdlp / baseline_sdlp
    srr_ratio = srr / baseline_srr
    
    # 综合评分（走神特征：速度波动↑，车道偏差↓，方向盘反转↓）
    score = (speed_var_ratio * 0.4 + 
             (1 - sdlp_ratio) * 0.3 + 
             (1 - srr_ratio) * 0.3)
    
    is_mw = score > 0.5
    
    return is_mw, score

class IMSMindWanderingModule:
    """IMS Mind Wandering 检测模块
    
    融合 EEG + 行为 + 眼动特征
    """
    
    def __init__(self):
        self.eeg_detector = MindWanderingDetector()
        self.behavior_weight = 0.3
        self.eeg_weight = 0.5
        self.eye_weight = 0.2
        
    def process_frame(self, frame, vehicle_data, eeg_data=None):
        """处理单帧数据
        
        Args:
            frame: 红外图像, shape=(H, W, C)
            vehicle_data: 车辆数据 dict
                - speed: 当前速度 km/h
                - lateral_offset: 横向偏移 m
                - steering_angle: 方向盘角度 deg
            eeg_data: EEG 数据（可选）, shape=(n_channels, n_samples)
        
        Returns:
            result: 检测结果 dict
                - is_mind_wandering: bool
                - confidence: float
                - alert_level: int (0-3)
        """
        scores = {}
        
        # 1. 行为特征评分
        behavior_score = self._analyze_behavior(vehicle_data)
        scores['behavior'] = behavior_score
        
        # 2. EEG 特征评分（如果可用）
        if eeg_data is not None:
            eeg_score = self._analyze_eeg(eeg_data)
            scores['eeg'] = eeg_score
        else:
            scores['eeg'] = 0.5  # 无数据时中性值
        
        # 3. 眼动特征评分
        eye_score = self._analyze_eye_movement(frame)
        scores['eye'] = eye_score
        
        # 4. 加权融合
        final_score = (
            scores['behavior'] * self.behavior_weight +
            scores['eeg'] * self.eeg_weight +
            scores['eye'] * self.eye_weight
        )
        
        # 5. 决策输出
        is_mw = final_score > 0.6
        alert_level = self._determine_alert_level(final_score)
        
        return {
            'is_mind_wandering': is_mw,
            'confidence': final_score,
            'alert_level': alert_level,
            'scores': scores
        }
    
    def _analyze_behavior(self, vehicle_data):
        """分析驾驶行为特征"""
        # 速度波动
        speed_var = np.std(vehicle_data['speed_history'])
        
        # 车道偏差
        lateral_std = np.std(vehicle_data['lateral_offset_history'])
        
        # 方向盘反转率
        steering_reversals = vehicle_data['steering_reversals']
        
        # 综合评分（示例阈值）
        score = 0.0
        if speed_var > 5:  # km/h std
            score += 0.3
        if lateral_std < 0.1:  # 过于平滑
            score += 0.4
        if steering_reversals < 10:  # 过少
            score += 0.3
            
        return min(score, 1.0)
    
    def _analyze_eeg(self, eeg_data):
        """分析 EEG 特征"""
        # 简化实现，实际需要完整处理流程
        alpha_power = self.eeg_detector.extract_alpha_power(eeg_data)
        # 基线比较...
        return 0.6  # 示例值
    
    def _analyze_eye_movement(self, frame):
        """分析眼动特征（基于视觉的 Mind Wandering 检测）"""
        # 眼动熵（Gaze Entropy）降低表示走神
        # 注视时长增加但眨眼率降低
        return 0.4  # 示例值
    
    def _determine_alert_level(self, score):
        """确定警告等级"""
        if score < 0.4:
            return 0  # 正常
        elif score < 0.6:
            return 1  # 提示
        elif score < 0.8:
            return 2  # 警告
        else:
            return 3  # 干预