import numpy as np
from scipy import signal as sig
from typing import Tuple

class fNIRSProcessor:
    """
    fNIRS 近红外光谱信号处理器
    用于提取前额叶血流动力学响应
    """
    
    def __init__(self, sample_rate: float = 10.0):
        """
        Args:
            sample_rate: 采样率 (Hz), fNIRS 通常 10-20 Hz
        """
        self.fs = sample_rate
        
        # fNIRS 滤波参数
        self.hbO_lowcut = 0.01  # Hz, 去除基线漂移
        self.hbO_highcut = 0.5  # Hz, 去除心跳（~1Hz）
        
    def preprocess(self, raw_hbO: np.ndarray, raw_hbR: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        fNIRS 预处理链路
        
        Args:
            raw_hbO: 原始氧合血红蛋白信号, shape=(n_samples, n_channels)
            raw_hbR: 原始脱氧血红蛋白信号, shape=(n_samples, n_channels)
            
        Returns:
            hbO_filtered, hbR_filtered: 滤波后信号
        """
        nyquist = self.fs / 2
        
        # 带通滤波 0.01-0.5 Hz
        b, a = sig.butter(
            4, 
            [self.hbO_lowcut/nyquist, self.hbO_highcut/nyquist], 
            btype='band'
        )
        
        hbO_filtered = sig.filtfilt(b, a, raw_hbO, axis=0)
        hbR_filtered = sig.filtfilt(b, a, raw_hbR, axis=0)
        
        # 运动伪影去除（基于方差的自适应阈值）
        for ch in range(hbO_filtered.shape[1]):
            # 计算滑动窗口方差
            window = 50  # 5秒 @ 10Hz
            var = np.convolve(
                np.var(hbO_filtered[:, ch:ch+1], axis=0), 
                np.ones(window)/window, 
                mode='same'
            )
            # 高方差段视为运动伪影，插值替换
            mask = var > np.median(var) * 3
            if np.any(mask):
                hbO_filtered[mask, ch] = np.interp(
                    np.where(mask)[0], 
                    np.where(~mask)[0], 
                    hbO_filtered[~mask, ch]
                )
        
        return hbO_filtered, hbR_filtered
    
    def extract_features(self, hbO: np.ndarray, hbR: np.ndarray) -> dict:
        """
        从 fNIRS 信号提取认知负荷特征
        
        Args:
            hbO: 预处理后 HbO 信号
            hbR: 预处理后 HbR 信号
            
        Returns:
            features: 认知负荷特征字典
        """
        features = {}
        
        # 1. HbO 均值（前额叶激活指标）
        features['hbO_mean'] = np.mean(hbO, axis=0)
        
        # 2. HbO 方差（认知负荷波动）
        features['hbO_std'] = np.std(hbO, axis=0)
        
        # 3. HbO/HbR 比值（氧合代谢率）
        features['hbO_hbR_ratio'] = np.mean(hbO) / (np.mean(hbR) + 1e-6)
        
        # 4. 前额叶左右不对称性（情绪/压力指标）
        n_channels = hbO.shape[1]
        left = np.mean(hbO[:, :n_channels//2])
        right = np.mean(hbO[:, n_channels//2:])
        features['frontal_asymmetry'] = left - right
        
        # 5. 频域特征（0.01-0.1Hz 为认知负荷频带）
        freqs = np.fft.rfftfreq(len(hbO), 1/self.fs)
        fft_hbO = np.abs(np.fft.rfft(hbO, axis=0))
        mask = (freqs >= 0.01) & (freqs <= 0.1)
        features['cognitive_band_power'] = np.mean(fft_hbO[mask], axis=0)
        
        return features


class TakeoverReadinessFusion:
    """
    多模态接管就绪性融合评估器
    
    论文核心：认知负荷×危险等级的交互效应
    """
    
    def __init__(self):
        # 上下文权重矩阵（论文核心发现）
        # [低危险, 高危险] × [低负荷, 高负荷]
        self.context_weights = {
            'fnirs': {
                ('low', 'low'): 0.3,   # 低危险低负荷：fNIRS 信号弱
                ('low', 'high'): 0.7,  # 低危险高负荷：fNIRS 信号强
                ('high', 'low'): 0.5,  # 高危险低负荷：中等
                ('high', 'high'): 0.8, # 高危险高负荷：最强
            },
            'driving': {
                ('low', 'low'): 0.7,   # 低危险低负荷：行为信号主导
                ('low', 'high'): 0.5,  
                ('high', 'low'): 0.8,  # 高危险低负荷：行为变化大
                ('high', 'high'): 0.6,
            },
            'eye': {
                ('low', 'low'): 0.5,
                ('low', 'high'): 0.6,
                ('high', 'low'): 0.7,
                ('high', 'high'): 0.5,
            }
        }
    
    def assess(self, fnirs_feat: dict, driving_feat: dict, eye_feat: dict,
               hazard_level: str, cognitive_load: str) -> dict:
        """
        评估接管就绪性
        
        Args:
            fnirs_feat: fNIRS 特征
            driving_feat: 驾驶行为特征
            eye_feat: 眼动特征
            hazard_level: 'low' or 'high'
            cognitive_load: 'low' or 'high'
            
        Returns:
            assessment: 就绪性评估结果
        """
        context = (hazard_level, cognitive_load)
        
        # 上下文自适应权重
        w_fnirs = self.context_weights['fnirs'][context]
        w_driving = self.context_weights['driving'][context]
        w_eye = self.context_weights['eye'][context]
        
        # 归一化
        total = w_fnirs + w_driving + w_eye
        w_fnirs /= total
        w_driving /= total
        w_eye /= total
        
        # 各模态就绪性分数（简化）
        fnirs_score = self._fnirs_readiness(fnirs_feat)
        driving_score = self._driving_readiness(driving_feat)
        eye_score = self._eye_readiness(eye_feat)
        
        # 融合
        fused_score = (
            w_fnirs * fnirs_score +
            w_driving * driving_score +
            w_eye * eye_score
        )
        
        return {
            'readiness_score': fused_score,
            'context': context,
            'fnirs_score': fnirs_score,
            'driving_score': driving_score,
            'eye_score': eye_score,
            'weights': {
                'fnirs': w_fnirs,
                'driving': w_driving,
                'eye': w_eye
            },
            'recommendation': 'ready' if fused_score > 0.6 else 'not_ready'
        }
    
    def _fnirs_readiness(self, feat: dict) -> float:
        """fNIRS 就绪性评分"""
        # 高 HbO + 低不对称 = 高就绪
        score = 0.5 + 0.3 * np.tanh(feat.get('hbO_mean', 0)) \
                    - 0.2 * np.tanh(abs(feat.get('frontal_asymmetry', 0)))
        return np.clip(score, 0, 1)
    
    def _driving_readiness(self, feat: dict) -> float:
        """驾驶行为就绪性评分"""
        score = 0.5 + 0.3 * feat.get('steering_activity', 0) \
                    + 0.2 * feat.get('speed_consistency', 0)
        return np.clip(score, 0, 1)
    
    def _eye_readiness(self, feat: dict) -> float:
        """眼动就绪性评分"""
        score = 0.5 + 0.3 * feat.get('gaze_road_ratio', 0) \
                    - 0.2 * feat.get('pupil_diameter_var', 0)
        return np.clip(score, 0, 1)


# 测试
if __name__ == "__main__":
    proc = fNIRSProcessor(sample_rate=10.0)
    
    # 模拟 fNIRS 数据（48通道，120秒）
    np.random.seed(42)
    raw_hbO = np.random.randn(1200, 48) * 0.1 + 0.5
    raw_hbR = np.random.randn(1200, 48) * 0.1 + 0.3
    
    hbO, hbR = proc.preprocess(raw_hbO, raw_hbR)
    features = proc.extract_features(hbO, hbR)
    
    print("fNIRS 特征:")
    for k, v in features.items():
        if isinstance(v, np.ndarray):
            print(f"  {k}: mean={v.mean():.4f}, std={v.std():.4f}")
        else:
            print(f"  {k}: {v:.4f}")
    
    # 融合评估
    fusion = TakeoverReadinessFusion()
    
    # 场景1: 高危险 + 高认知负荷
    result = fusion.assess(
        fnirs_feat=features,
        driving_feat={'steering_activity': 0.7, 'speed_consistency': 0.6},
        eye_feat={'gaze_road_ratio': 0.8, 'pupil_diameter_var': 0.3},
        hazard_level='high',
        cognitive_load='high'
    )
    print(f"\n场景: 高危险 + 高认知负荷")
    print(f"  就绪性分数: {result['readiness_score']:.3f}")
    print(f"  权重分配: fNIRS={result['weights']['fnirs']:.2f}, "
          f"驾驶={result['weights']['driving']:.2f}, "
          f"眼动={result['weights']['eye']:.2f}")
    print(f"  建议: {result['recommendation']}")

class VisionBasedCognitiveLoadEstimator:
    """
    基于论文发现的无 fNIRS 认知负荷推断器
    用眼动+驾驶行为替代脑信号
    """
    
    def __init__(self):
        # 从论文迁移的上下文权重（去掉 fNIRS 后重新分配）
        # 低危险时行为信号权重高，高危险时眼动权重高
        self.weights = {
            ('low', 'low'): {'driving': 0.6, 'eye': 0.4},
            ('low', 'high'): {'driving': 0.4, 'eye': 0.6},
            ('high', 'low'): {'driving': 0.7, 'eye': 0.3},
            ('high', 'high'): {'driving': 0.3, 'eye': 0.7},
        }
    
    def estimate(self, driving_feat: dict, eye_feat: dict, 
                 hazard_level: str) -> dict:
        """
        估算认知负荷等级
        
        Args:
            driving_feat: 方向盘/油门/刹车特征
            eye_feat: 眼动特征（注视、瞳孔）
            hazard_level: 当前危险等级
            
        Returns:
            result: 认知负荷评估
        """
        # 瞳孔直径是认知负荷的可靠 proxy
        pupil_load = self._pupil_based_load(eye_feat)
        steering_load = self._steering_based_load(driving_feat)
        
        # 简单规则推断负荷等级
        raw_load = 0.5 * pupil_load + 0.5 * steering_load
        cognitive_load = 'high' if raw_load > 0.6 else 'low'
        
        # 上下文权重
        context = (hazard_level, cognitive_load)
        w = self.weights.get(context, {'driving': 0.5, 'eye': 0.5})
        
        # 就绪性评估
        readiness = (
            w['driving'] * (1 - steering_load) + 
            w['eye'] * eye_feat.get('gaze_road_ratio', 0.5)
        )
        
        return {
            'cognitive_load': cognitive_load,
            'raw_load_score': raw_load,
            'readiness': readiness,
            'recommendation': 'alert' if readiness < 0.5 else 'normal'
        }
    
    def _pupil_based_load(self, eye_feat: dict) -> float:
        """瞳孔直径推断认知负荷"""
        # 认知负荷↑ → 瞳孔直径↑（0.1-0.5mm 变化）
        return np.clip(eye_feat.get('pupil_diameter_var', 0.3), 0, 1)
    
    def _steering_based_load(self, driving_feat: dict) -> float:
        """方向盘微修正频率推断负荷"""
        # 认知负荷↑ → 方向盘微修正↓（注意力被占用）
        corrections = driving_feat.get('steering_corrections_per_min', 30)
        return np.clip(1 - corrections / 60, 0, 1)


# 测试
if __name__ == "__main__":
    est = VisionBasedCognitiveLoadEstimator()
    
    # 场景: 高速公路，高危险
    result = est.estimate(
        driving_feat={'steering_corrections_per_min': 15},
        eye_feat={'pupil_diameter_var': 0.6, 'gaze_road_ratio': 0.7},
        hazard_level='high'
    )
    print(f"认知负荷: {result['cognitive_load']}")
    print(f"就绪性: {result['readiness']:.3f}")
    print(f"建议: {result['recommendation']}")