import numpy as np
from scipy.signal import welch
from sklearn.linear_model import Lasso
from sklearn.preprocessing import StandardScaler

# EEG采集参数（论文复现）
EEG_CHANNELS = ['Fp1','Fp2','F3','F4','Fz','C3','C4','Cz','P3','P4','Pz','O1','O2','T5','T6','F7','F8']
SAMPLING_RATE = 500  # Hz
BANDS = {
    'delta': (1, 4),
    'theta': (4, 8),
    'alpha': (8, 13),
    'low_beta': (13, 20),
    'high_beta': (20, 30),
    'gamma': (30, 45)
}

def extract_eeg_features(eeg_data, fs=500):
    """
    从EEG信号中提取频谱特征
    
    Args:
        eeg_data: shape=(n_channels, n_samples)
        fs: 采样率
        
    Returns:
        features: dict, {channel_band: power_value}
    """
    features = {}
    for ch_idx, ch_name in enumerate(EEG_CHANNELS):
        # Welch法估计功率谱密度
        freqs, psd = welch(eeg_data[ch_idx], fs=fs, nperseg=fs*2)
        
        for band_name, (fmin, fmax) in BANDS.items():
            # 计算频段平均功率
            idx = np.logical_and(freqs >= fmin, freqs <= fmax)
            band_power = np.mean(psd[idx])
            features[f"{ch_name}_{band_name}"] = band_power
    
    return features


def select_eeg_features_lasso(X, y, n_select=2):
    """
    使用L1正则化(LASSO)选择最具判别力的EEG特征
    
    Args:
        X: shape=(n_samples, n_features), EEG特征矩阵
        y: shape=(n_samples,), 标签 (0=健康, 1=MCI)
        n_select: 选择的特征数
        
    Returns:
        selected_indices: 选中的特征索引
        selected_names: 选中的特征名
    """
    # 标准化
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    
    # LASSO特征选择
    lasso = Lasso(alpha=0.01, max_iter=10000)
    lasso.fit(X_scaled, y)
    
    # 按系数绝对值排序
    importance = np.abs(lasso.coef_)
    selected_indices = np.argsort(importance)[-n_select:]
    
    return selected_indices, importance

def extract_gaze_features(gaze_x, gaze_y, sample_rate=1000):
    """
    从眼动数据中提取注视稳定性特征
    
    Args:
        gaze_x: shape=(n_samples,), X坐标（度）
        gaze_y: shape=(n_samples,), Y坐标（度）
        sample_rate: 采样率
        
    Returns:
        features: dict, 眼动变异性特征
    """
    # 排除扫视区间（仅保留注视数据）
    # 实际实现需使用IVMC或Eyelink自带扫视检测
    
    # 计算注视点位置标准差
    std_x = np.std(gaze_x)
    std_y = np.std(gaze_y)
    
    features = {
        'R_std_x': std_x,  # X方向注视稳定性
        'L_std_x': std_x,  # 论文中使用的命名
        'R_std_y': std_y,
        'L_std_y': std_y,
        'gaze_dispersion': np.sqrt(std_x**2 + std_y**2),
        'gaze_area': np.pi * std_x * std_y  # 注视椭圆面积
    }
    
    return features

from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import LeaveOneOut
from sklearn.metrics import roc_auc_score

def fused_cognitive_classifier(eeg_features, gaze_features, labels):
    """
    EEG+眼动融合认知状态分类器
    
    论文方法复现：
    1. 模型(a): 全EEG特征
    2. 模型(b): LASSO选择的EEG特征
    3. 模型(c): LASSO-EEG + 眼动特征
    
    Args:
        eeg_features: shape=(n_samples, n_eeg_features)
        gaze_features: shape=(n_samples, n_gaze_features)
        labels: shape=(n_samples,)
        
    Returns:
        results: dict, 三个模型的AUC
    """
    loocv = LeaveOneOut()
    
    # 模型(a): 全EEG特征
    model_a = LogisticRegression(class_weight='balanced', max_iter=5000)
    pred_a = _loocv_predict(model_a, eeg_features, labels, loocv)
    auc_a = roc_auc_score(labels, pred_a)
    
    # 模型(b): LASSO选择后的EEG特征
    selected_idx, _ = select_eeg_features_lasso(eeg_features, labels, n_select=2)
    eeg_selected = eeg_features[:, selected_idx]
    model_b = LogisticRegression(class_weight='balanced', max_iter=5000)
    pred_b = _loocv_predict(model_b, eeg_selected, labels, loocv)
    auc_b = roc_auc_score(labels, pred_b)
    
    # 模型(c): LASSO-EEG + 眼动特征
    fused = np.hstack([eeg_selected, gaze_features])
    model_c = LogisticRegression(class_weight='balanced', max_iter=5000)
    pred_c = _loocv_predict(model_c, fused, labels, loocv)
    auc_c = roc_auc_score(labels, pred_c)
    
    return {
        'model_a_all_eeg': auc_a,      # 0.52
        'model_b_lasso_eeg': auc_b,    # 0.64
        'model_c_fused': auc_c          # 0.78
    }


def _loocv_predict(model, X, y, loocv):
    """留一交叉验证预测"""
    preds = np.zeros(len(y))
    for train_idx, test_idx in loocv.split(X):
        model.fit(X[train_idx], y[train_idx])
        preds[test_idx] = model.predict_proba(X[test_idx])[:, 1]
    return preds

def cabin_gaze_features(gaze_data, fps=30):
    """
    座舱场景眼动特征提取（增强版）
    
    相比论文的2维特征，座舱场景可提取更丰富的眼动指标
    """
    features = {}
    
    # 1. 注视稳定性（论文方法）
    features['fixation_std_x'] = np.std(gaze_data['x'])
    features['fixation_std_y'] = np.std(gaze_data['y'])
    
    # 2. PERCLOS（闭眼比例）
    ear = gaze_data['ear']  # 眼睛纵横比
    features['perclos_80'] = np.mean(ear < 0.2) * 100  # 80%闭眼阈值
    
    # 3. 扫视熵（视线方向分布的熵）
    gaze_bins = np.histogram2d(gaze_data['x'], gaze_data['y'], 
                                bins=9, range=[[-30,30],[-30,30]])[0]
    gaze_probs = gaze_bins / gaze_bins.sum() + 1e-10
    features['gaze_entropy'] = -np.sum(gaze_probs * np.log(gaze_probs))
    
    # 4. 视线偏离道路时间比
    features['off_road_ratio'] = np.mean(
        (np.abs(gaze_data['x']) > 15) | (np.abs(gaze_data['y']) > 10)
    ) * 100
    
    # 5. 微扫视频率（认知负荷指标）
    saccade_vel = np.sqrt(np.diff(gaze_data['x'])**2 + np.diff(gaze_data['y'])**2)
    features['microsaccade_rate'] = np.mean(saccade_vel > 50)  # >50°/s视为微扫视
    
    return features

"""
座舱认知分心检测改进方案

改进1: 滑动窗口实时检测
"""
class RealTimeCognitiveDetector:
    def __init__(self, window_sec=4, fps=30):
        self.window = int(window_sec * fps)
        self.model = LogisticRegression(class_weight='balanced')
        self.buffer = []
        
    def update(self, gaze_features, eeg_features=None):
        """实时更新认知状态评估"""
        self.buffer.append((gaze_features, eeg_features))
        if len(self.buffer) > self.window:
            self.buffer.pop(0)
        
        if len(self.buffer) >= self.window // 2:  # 至少2s数据
            features = self._aggregate()
            prob = self.model.predict_proba([features])[0, 1]
            return prob  # 认知分心概率
        return 0.0
    
    def _aggregate(self):
        """聚合窗口内特征"""
        gaze_arr = [f[0] for f in self.buffer if f[0] is not None]
        eeg_arr = [f[1] for f in self.buffer if f[1] is not None]
        
        agg = {}
        if gaze_arr:
            agg['gaze_std'] = np.std([g['fixation_std_x'] for g in gaze_arr])
            agg['gaze_entropy_mean'] = np.mean([g['gaze_entropy'] for g in gaze_arr])
        if eeg_arr:
            agg['eeg_delta'] = np.mean([e['O1_delta'] for e in eeg_arr])
            agg['eeg_beta'] = np.mean([e['T6_low_beta'] for e in eeg_arr])
        
        return list(agg.values())


"""
改进2: 增强眼动特征集（替代EEG）
在无EEG传感器时，用更丰富的眼动特征达到接近的性能
"""
ENHANCED_GAZE_FEATURES = [
    'fixation_std_x',          # 论文原始特征
    'fixation_std_y',          # 论文原始特征
    'perclos_80',              # PERCLOS
    'gaze_entropy',            # 视线熵
    'off_road_ratio',          # 视线偏离比
    'microsaccade_rate',       # 微扫视频率
    'blink_rate',              # 眨眼频率
    'pupil_diameter_cv',       # 瞳孔直径变异系数
    'saccade_peak_velocity',   # 扫视峰值速度
    'fixation_duration_cv',    # 注视持续时间变异系数
]