"""
可穿戴fNIRS认知疲劳检测信号处理框架

复现论文核心信号处理流程 + 扩展车载应用

依赖: numpy>=1.24, scipy>=1.11, pywt>=1.5
"""

import numpy as np
from scipy.signal import butter, filtfilt, detrend
from scipy.stats import pearsonr
import pywt
from typing import Tuple, Dict


class WearableFNIRSProcessor:
    """
    可穿戴fNIRS信号处理器
    
    核心功能:
    1. 原始光强 → 血氧变化（Modified Beer-Lambert Law）
    2. 运动伪影去除（小波域滤波）
    3. 认知疲劳指标提取（VLF漂移 + oxy-Hb趋势）
    4. 实时分类（正常/轻度疲劳/重度疲劳）
    """
    
    def __init__(self, config: dict):
        # 光学参数
        self.wavelengths = config.get('wavelengths', [760, 850])  # nm
        self.dpf = config.get('dpf', 6.0)  # 差分光程因子
        self.source_detector_dist = config.get('distance', 0.015)  # 15mm
        
        # 滤波参数
        self.low_cut = 0.01  # Hz, 0.01 Hz (VLF)
        self.high_cut = 0.5  # Hz
        self.fs = config.get('sampling_rate', 10)  # Hz
        
        # 疲劳指标参数
        self.window_sec = config.get('window_sec', 60)  # 60秒滑动窗口
        self.fatigue_threshold = config.get('fatigue_threshold', 0.3)
        
    def raw_to_hb(
        self, 
        raw_760: np.ndarray, 
        raw_850: np.ndarray
    ) -> Dict[str, np.ndarray]:
        """
        将原始光强转换为血氧浓度变化
        
        基于Modified Beer-Lambert Law (MBLL)
        
        Args:
            raw_760: 760nm波长光强序列
            raw_850: 850nm波长光强序列
            
        Returns:
            dict: oxy_Hb, deoxy_Hb, total_Hb 变化量 (ΔμM)
        """
        # 光密度变化
        od_760 = -np.log10(raw_760 / np.mean(raw_760[:100]))
        od_850 = -np.log10(raw_850 / np.mean(raw_850[:100]))
        
        # 消光系数 (cm^-1 / M)
        # 760nm: oxy=0.484, deoxy=1.492
        # 850nm: oxy=1.026, deoxy=0.555
        eps_760 = np.array([0.484, 1.492])
        eps_850 = np.array([1.026, 0.555])
        
        # MBLL矩阵
        # ΔOD = ε * Δc * d * DPF
        d = self.source_detector_dist * 100  # 转cm
        DPF = self.dpf
        
        # 线性方程组求解
        # [ΔOD_760]   [ε_760_oxy  ε_760_deoxy] [Δc_oxy]
        # [ΔOD_850] = [ε_850_oxy  ε_850_deoxy] [Δc_deoxy] * d * DPF
        eps_matrix = np.array([
            [eps_760[0], eps_760[1]],
            [eps_850[0], eps_850[1]]
        ])
        
        # 解线性方程
        od_vector = np.array([od_760, od_850])  # (2, N)
        delta_c = np.linalg.solve(
            eps_matrix * d * DPF,
            od_vector
        )
        
        return {
            'oxy_Hb': delta_c[0],
            'deoxy_Hb': delta_c[1],
            'total_Hb': delta_c[0] + delta_c[1]
        }
    
    def remove_motion_artifact(
        self, 
        signal: np.ndarray,
        wavelet: str = 'db4',
        threshold: float = 2.0
    ) -> np.ndarray:
        """
        基于小波变换的运动伪影去除
        
        论文方法：高频小波系数阈值化
        """
        # 小波分解
        coeffs = pywt.wavedec(signal, wavelet, level=6)
        
        # 阈值化高频系数（去运动伪影）
        for i in range(1, len(coeffs)):
            sigma = np.std(coeffs[i])
            adaptive_threshold = threshold * sigma
            coeffs[i] = pywt.threshold(
                coeffs[i], 
                adaptive_threshold, 
                mode='soft'
            )
        
        # 重建信号
        cleaned = pywt.waverec(coeffs, wavelet)
        return cleaned[:len(signal)]
    
    def bandpass_filter(self, signal: np.ndarray) -> np.ndarray:
        """
        带通滤波 0.01-0.5 Hz
        
        保留血流动力学响应频段
        去除心跳(~1Hz)和呼吸(~0.2Hz)干扰
        """
        nyquist = self.fs / 2
        low = self.low_cut / nyquist
        high = self.high_cut / nyquist
        
        b, a = butter(4, [low, high], btype='band')
        return filtfilt(b, a, signal)
    
    def extract_fatigue_features(
        self, 
        oxy_hb: np.ndarray,
        window_sec: int = 60
    ) -> Dict[str, float]:
        """
        提取认知疲劳特征
        
        论文发现：
        - VLF (<0.02Hz) 漂移与疲劳发展相关
        - oxy-Hb 在高认知负荷下先升后降
        - 疲劳期VLF漂移持续到休息阶段
        
        Returns:
            dict: 疲劳相关特征
        """
        window_samples = int(window_sec * self.fs)
        features = {}
        
        # 1. VLF漂移 (< 0.02 Hz)
        vlf_signal = self._vlf_filter(oxy_hb)
        features['vlf_drift'] = np.mean(vlf_signal[-window_samples:])
        features['vlf_slope'] = np.polyfit(
            np.arange(window_samples), 
            vlf_signal[-window_samples:], 1
        )[0]
        
        # 2. oxy-Hb均值变化
        features['oxy_mean'] = np.mean(oxy_hb[-window_samples:])
        features['oxy_std'] = np.std(oxy_hb[-window_samples:])
        
        # 3. 氧合指数 (oxy_Hb / total_Hb)
        total = np.abs(oxy_hb[-window_samples:]).sum() + 1e-6
        features['oxygenation_index'] = np.mean(oxy_hb[-window_samples:]) / total
        
        # 4. 认知负荷波动
        features['oxy_range'] = np.ptp(oxy_hb[-window_samples:])
        
        return features
    
    def _vlf_filter(self, signal: np.ndarray) -> np.ndarray:
        """超低频滤波 (< 0.02 Hz)"""
        nyquist = self.fs / 2
        low = 0.002 / nyquist
        high = 0.02 / nyquist
        b, a = butter(2, [low, high], btype='band')
        return filtfilt(b, a, signal)
    
    def classify_fatigue(
        self, 
        features: Dict[str, float]
    ) -> Tuple[str, float]:
        """
        认知疲劳分类
        
        基于论文发现：
        - VLF漂移方向和幅度区分疲劳程度
        - oxy-Hb从上升转为下降标志疲劳临界点
        
        Returns:
            level: "alert" / "mild_fatigue" / "severe_fatigue"
            confidence: 0-1
        """
        vlf = features['vlf_drift']
        slope = features['vlf_slope']
        
        # 简化分类逻辑（实际应使用训练好的ML模型）
        if slope > 0.001 and vlf > 0:
            return "alert", 0.8
        elif -0.005 < slope <= 0.001:
            return "mild_fatigue", 0.7
        else:
            return "severe_fatigue", 0.75
    
    def process_stream(
        self, 
        raw_760: np.ndarray, 
        raw_850: np.ndarray
    ) -> Dict:
        """
        完整信号处理管线
        
        输入: 原始双波长光强序列
        输出: 疲劳评估结果
        """
        # 1. 原始→血氧
        hb = self.raw_to_hb(raw_760, raw_850)
        
        # 2. 运动伪影去除
        oxy_clean = self.remove_motion_artifact(hb['oxy_Hb'])
        
        # 3. 带通滤波
        oxy_filtered = self.bandpass_filter(oxy_clean)
        
        # 4. 特征提取
        features = self.extract_fatigue_features(oxy_filtered)
        
        # 5. 分类
        level, confidence = self.classify_fatigue(features)
        
        return {
            'fatigue_level': level,
            'confidence': confidence,
            'features': features,
            'oxy_hb_filtered': oxy_filtered
        }


# 测试
if __name__ == "__main__":
    # 模拟fNIRS数据
    fs = 10  # 10 Hz
    duration = 300  # 5分钟
    t = np.arange(0, duration, 1/fs)
    
    # 模拟正常状态→认知负荷→疲劳的血氧变化
    # 正常阶段 (0-120s): 稳定波动
    normal = 0.1 * np.sin(2 * np.pi * 0.05 * t[:120*fs])
    
    # 认知负荷阶段 (120-240s): oxy-Hb上升
    load = 0.3 * np.sin(2 * np.pi * 0.05 * t[120*fs:240*fs]) + 0.5 * np.linspace(0, 1, 120*fs)
    
    # 疲劳阶段 (240-300s): oxy-Hb下降 + VLF漂移
    fatigue = -0.2 * np.sin(2 * np.pi * 0.05 * t[240*fs:]) + np.linspace(0.5, -0.3, 60*fs)
    
    oxy_hb_true = np.concatenate([normal, load, fatigue])
    
    # 添加噪声和运动伪影
    noise = np.random.normal(0, 0.05, len(t))
    motion = np.zeros_like(t)
    motion[150*fs:155*fs] = 2.0  # 突发运动伪影
    
    raw_850 = 10**(-oxy_hb_true + noise + motion) * 1e-3
    raw_760 = 10**(-0.5 * oxy_hb_true + noise + motion) * 1e-3
    
    # 处理
    processor = WearableFNIRSProcessor({'sampling_rate': fs})
    result = processor.process_stream(raw_760, raw_850)
    
    print(f"疲劳等级: {result['fatigue_level']}")
    print(f"置信度: {result['confidence']:.2f}")
    print(f"VLF漂移: {result['features']['vlf_drift']:.4f}")
    print(f"VLF斜率: {result['features']['vlf_slope']:.6f}")