import numpy as np
from scipy import signal
from scipy.fft import fft
from typing import Tuple

class SSVEPFatiguePredictor:
    """
    SSVEP（稳态视觉诱发电位）疲劳预测器
    
    论文发现: SSVEP可在注意力衰减前5秒预测错误
    方法: 对仪表盘特定频率闪烁 stimulus → EEG检测同频率响应
    
    航空应用: 飞行员仪表扫描质量实时评估
    汽车迁移: 可用于HUD闪烁标记→驾驶员注意力检测
    """
    
    def __init__(self, fs=500, target_freqs=[10, 15, 20]):
        """
        Args:
            fs: EEG采样率 (Hz)
            target_freqs: 目标闪烁频率（仪表盘stimulus频率）
        """
        self.fs = fs
        self.target_freqs = target_freqs
        # 5秒预测窗口
        self.prediction_window_s = 5
        self.window_samples = int(self.prediction_window_s * fs)
        
    def extract_ssvep(self, eeg_signal: np.ndarray, 
                       channel_idx: int = 0) -> dict:
        """
        提取SSVEP响应
        
        Args:
            eeg_signal: EEG信号, shape=(channels, samples)
            channel_idx: 分析的通道索引
            
        Returns:
            dict: 各目标频率的SSVEP信噪比
        """
        if eeg_signal.ndim == 1:
            eeg_signal = eeg_signal[np.newaxis, :]
            
        signal_data = eeg_signal[channel_idx]
        
        # FFT分析
        freqs = np.fft.rfftfreq(len(signal_data), 1/self.fs)
        spectrum = np.abs(np.fft.rfft(signal_data))
        
        results = {}
        for target in self.target_freqs:
            # 找到目标频率附近的bin
            target_bin = np.argmin(np.abs(freqs - target))
            
            # 信号功率（目标频率±0.5Hz）
            signal_mask = (freqs >= target - 0.5) & (freqs <= target + 0.5)
            signal_power = np.sum(spectrum[signal_mask]**2)
            
            # 噪声功率（周围频段）
            noise_mask = ((freqs >= target - 3) & (freqs < target - 1)) | \
                         ((freqs > target + 1) & (freqs <= target + 3))
            noise_power = np.mean(spectrum[noise_mask]**2)
            
            # SNR
            snr = 10 * np.log10(signal_power / (noise_power + 1e-10))
            
            results[f'{target}Hz'] = {
                'snr_db': snr,
                'signal_power': signal_power,
                'noise_power': noise_power,
                'amplitude': np.sqrt(signal_power)
            }
        
        return results
    
    def predict_attention_lapse(self, eeg_window: np.ndarray) -> dict:
        """
        预测5秒内是否会发生注意力衰减
        
        论文发现: SSVEP响应在注意力衰减前5秒显著降低
        
        Args:
            eeg_window: 5秒EEG窗口, shape=(channels, samples)
            
        Returns:
            dict: 预测结果
        """
        ssvep = self.extract_ssvep(eeg_window)
        
        # 所有目标频率的平均SNR
        avg_snr = np.mean([v['snr_db'] for v in ssvep.values()])
        
        # 阈值判定（论文基线）
        SNR_THRESHOLD = 3.0  # dB
        risk_score = max(0, min(1, (SNR_THRESHOLD - avg_snr) / SNR_THRESHOLD))
        
        return {
            'avg_snr_db': avg_snr,
            'risk_score': risk_score,
            'prediction': 'high_risk' if risk_score > 0.6 else 
                         'moderate' if risk_score > 0.3 else 'low',
            'time_to_lapse_s': self.prediction_window_s if risk_score > 0.6 else None,
            'ssvep_details': ssvep
        }


# 测试
if __name__ == "__main__":
    np.random.seed(42)
    predictor = SSVEPFatiguePredictor(fs=500)
    
    # 模拟5秒EEG数据（500Hz × 5秒 = 2500样本）
    t = np.arange(2500) / 500
    
    # 正常状态：SSVEP响应强（10Hz闪烁 + 噪声）
    normal_eeg = (5 * np.sin(2*np.pi*10*t) +    # 10Hz SSVEP
                  3 * np.sin(2*np.pi*15*t) +     # 15Hz SSVEP
                  np.random.normal(0, 1, 2500))   # 噪声
    
    # 疲劳状态：SSVEP响应弱
    fatigued_eeg = (1.5 * np.sin(2*np.pi*10*t) +  # 幅度降低70%
                    0.8 * np.sin(2*np.pi*15*t) +
                    np.random.normal(0, 1.5, 2500))
    
    print("=== 正常状态 ===")
    result_normal = predictor.predict_attention_lapse(normal_eeg[np.newaxis, :])
    print(f"平均SNR: {result_normal['avg_snr_db']:.2f} dB")
    print(f"风险评分: {result_normal['risk_score']:.2f}")
    print(f"预测: {result_normal['prediction']}")
    
    print("\n=== 疲劳状态 ===")
    result_fatigued = predictor.predict_attention_lapse(fatigued_eeg[np.newaxis, :])
    print(f"平均SNR: {result_fatigued['avg_snr_db']:.2f} dB")
    print(f"风险评分: {result_fatigued['risk_score']:.2f}")
    print(f"预测: {result_fatigued['prediction']}")
    if result_fatigued.get('time_to_lapse_s'):
        print(f"预计{result_fatigued['time_to_lapse_s']}秒内注意力衰减")

# 飞行员经验与扫描模式对比
EXPERIENCE_DATA = {
    'novice': {
        'flight_hours': '40-70h',
        'aoi_dwell_ms': 274,
        'scan_transitions': 'few, repetitive',
        'pattern': 'unstructured',
        'pupil_change': 'high variance'
    },
    'experienced': {
        'flight_hours': '1500-2150h',
        'aoi_dwell_ms': 200,
        'scan_transitions': 'moderate, varied',
        'pattern': 'semi-structured',
        'pupil_change': 'moderate variance'
    },
    'expert': {
        'flight_hours': '3000-10300h',
        'aoi_dwell_ms': 167,
        'scan_transitions': 'many, comprehensive',
        'pattern': 'structured + adaptive',
        'pupil_change': 'low variance (stable)'
    }
}

# 关键发现: 84%重大航空事故与监控不足有关 (NTSB数据)
# 专家飞行员AOI停留时间比新手短39%