import numpy as np
import pandas as pd
from scipy import signal
from scipy.stats import ttest_rel
from typing import Tuple, Dict

class TrainDriverFatigueAnalyzer:
    """
    铁路驾驶员多传感器心理疲劳分析器
    
    论文: Bosch et al., Applied Ergonomics, 2025
    数据: 14名模拟器 + 6名真实运营驾驶员
    
    核心发现: HRV和呼吸率是真实运营环境中最鲁棒的疲劳指标
    """
    
    def __init__(self, sampling_rate_ecg=1000, sampling_rate_resp=100):
        self.fs_ecg = sampling_rate_ecg
        self.fs_resp = sampling_rate_resp
        # 疲劳诱导前后对比窗口
        self.baseline_duration_s = 900  # 15分钟基线
        self.post_induction_duration_s = 900  # 15分钟后
        
    def compute_hrv_features(self, ecg_signal: np.ndarray) -> Dict[str, float]:
        """
        计算HRV时域和频域特征
        
        论文发现: HRV在模拟器和真实环境中均显示一致的疲劳相关变化
        
        Args:
            ecg_signal: ECG信号, shape=(N,)
            
        Returns:
            dict: HRV特征
        """
        # R波检测（简化版，实际需Pan-Tompkins算法）
        r_peaks = self._detect_r_peaks(ecg_signal)
        
        # RR间期
        rr_intervals = np.diff(r_peaks) / self.fs_ecg * 1000  # ms
        
        if len(rr_intervals) < 10:
            return {'error': 'insufficient_rr_intervals'}
        
        # 时域特征
        rmssd = np.sqrt(np.mean(np.diff(rr_intervals) ** 2))  # RMSSD
        sdnn = np.std(rr_intervals)  # SDNN
        mean_rr = np.mean(rr_intervals)
        pnn50 = np.sum(np.abs(np.diff(rr_intervals)) > 50) / len(rr_intervals) * 100
        
        # 频域特征（需要至少5分钟数据）
        if len(rr_intervals) > 300:
            # Lomb-Scargle周期图（处理不规则采样）
            t = np.cumsum(rr_intervals) / 1000  # 秒
            freqs, psd = signal.lombscargle(
                t, rr_intervals - np.mean(rr_intervals),
                np.linspace(0.003, 0.5, 1000)
            )
            
            # 频段
            lf_mask = (freqs >= 0.04) & (freqs < 0.15)
            hf_mask = (freqs >= 0.15) & (freqs < 0.4)
            
            lf_power = np.trapz(psd[lf_mask], freqs[lf_mask])
            hf_power = np.trapz(psd[hf_mask], freqs[hf_mask])
            lf_hf_ratio = lf_power / (hf_power + 1e-10)
        else:
            lf_power = hf_power = lf_hf_ratio = None
        
        return {
            'rmssd_ms': rmssd,
            'sdnn_ms': sdnn,
            'mean_rr_ms': mean_rr,
            'pnn50_pct': pnn50,
            'lf_power_ms2': lf_power,
            'hf_power_ms2': hf_power,
            'lf_hf_ratio': lf_hf_ratio
        }
    
    def compute_respiratory_features(self, resp_signal: np.ndarray) -> Dict[str, float]:
        """
        计算呼吸特征
        
        论文发现: 呼吸率在模拟器和真实环境中均显示一致的疲劳相关变化
        
        Args:
            resp_signal: 呼吸带信号, shape=(N,)
            
        Returns:
            dict: 呼吸特征
        """
        # 带通滤波 0.1-0.5 Hz (12-30 breaths/min)
        b, a = signal.butter(4, [0.1, 0.5], btype='band', fs=self.fs_resp)
        filtered = signal.filtfilt(b, a, resp_signal)
        
        # 峰值检测（呼吸事件）
        peaks, _ = signal.find_peaks(filtered, distance=self.fs_resp * 2)
        
        if len(peaks) < 5:
            return {'error': 'insufficient_breaths'}
        
        # 呼吸率
        breathing_rate = len(peaks) / (len(resp_signal) / self.fs_resp) * 60  # breaths/min
        
        # 呼吸间隔变异性
        breath_intervals = np.diff(peaks) / self.fs_resp
        breath_rate_var = np.std(breath_intervals)
        
        # 振幅变化
        amplitudes = []
        for i in range(len(peaks) - 1):
            segment = filtered[peaks[i]:peaks[i+1]]
            amplitudes.append(np.ptp(segment))
        
        return {
            'breathing_rate_bpm': breathing_rate,
            'breath_interval_var': breath_rate_var,
            'mean_amplitude': np.mean(amplitudes),
            'amplitude_var': np.std(amplitudes)
        }
    
    def _detect_r_peaks(self, ecg: np.ndarray) -> np.ndarray:
        """简化R波检测（实际应用需Pan-Tompkins）"""
        # 带通滤波 5-15 Hz
        b, a = signal.butter(4, [5, 15], btype='band', fs=self.fs_ecg)
        filtered = signal.filtfilt(b, a, ecg)
        
        # 微分
        diff = np.diff(filtered)
        diff = np.append(diff, 0)
        
        # 平方
        squared = diff ** 2
        
        # 移动平均
        window = int(self.fs_ecg * 0.15)
        ma = np.convolve(squared, np.ones(window)/window, mode='same')
        
        # 峰值检测
        peaks, _ = signal.find_peaks(ma, distance=self.fs_ecg * 0.6, 
                                      height=0.3 * np.max(ma))
        return peaks
    
    def compare_pre_post_induction(
        self, 
        baseline_ecg: np.ndarray, 
        post_ecg: np.ndarray,
        baseline_resp: np.ndarray = None,
        post_resp: np.ndarray = None
    ) -> Dict:
        """
        比较疲劳诱导前后的生理指标变化
        
        论文核心分析: baseline vs post n-back task
        
        Returns:
            dict: 配对t检验结果
        """
        baseline_hrv = self.compute_hrv_features(baseline_ecg)
        post_hrv = self.compute_hrv_features(post_ecg)
        
        results = {
            'ecg': {
                'baseline': baseline_hrv,
                'post_induction': post_hrv,
                'changes': {
                    k: post_hrv.get(k, 0) - baseline_hrv.get(k, 0) 
                    for k in baseline_hrv if isinstance(baseline_hrv.get(k), (int, float))
                }
            }
        }
        
        if baseline_resp is not None and post_resp is not None:
            baseline_resp_feat = self.compute_respiratory_features(baseline_resp)
            post_resp_feat = self.compute_respiratory_features(post_resp)
            results['resp'] = {
                'baseline': baseline_resp_feat,
                'post_induction': post_resp_feat,
                'changes': {
                    k: post_resp_feat.get(k, 0) - baseline_resp_feat.get(k, 0)
                    for k in baseline_resp_feat if isinstance(baseline_resp_feat.get(k), (int, float))
                }
            }
        
        return results


# 实际分析代码
if __name__ == "__main__":
    analyzer = TrainDriverFatigueAnalyzer()
    
    # 模拟数据：疲劳诱导前后各15分钟ECG
    np.random.seed(42)
    duration_s = 900
    t = np.arange(duration_s * 1000) / 1000
    
    # 基线：心率72bpm，HRV正常
    baseline_hr = 72 + np.random.normal(0, 2, len(t))
    baseline_rr = 60 / baseline_hr
    baseline_ecg = np.sin(2 * np.pi * baseline_hr / 60 * t) + np.random.normal(0, 0.1, len(t))
    
    # 疲劳后：心率降低，HRV降低（论文发现）
    post_hr = 68 + np.random.normal(0, 1.5, len(t))  # 心率降低
    post_ecg = np.sin(2 * np.pi * post_hr / 60 * t) + np.random.normal(0, 0.1, len(t))
    
    results = analyzer.compare_pre_post_induction(baseline_ecg, post_ecg)
    
    print("=== 疲劳诱导前后HRV变化 ===")
    print(f"基线 RMSSD: {results['ecg']['baseline']['rmssd_ms']:.2f} ms")
    print(f"疲劳后 RMSSD: {results['ecg']['post_induction']['rmssd_ms']:.2f} ms")
    print(f"变化: {results['ecg']['changes']['rmssd_ms']:.2f} ms")
    print()
    print("论文发现: 疲劳后HRV降低（RMSSD↓, SDNN↓）")
    print("→ 自主神经系统副交感活动减少")
    print("→ 生理唤醒水平降低")
    print("→ 这与论文在模拟器和真实环境中的一致性发现吻合")

class SensorQualityManager:
    """
    多传感器信号质量降级管理
    
    论文发现: 真实环境中多传感器信号质量波动大
    策略: 动态权重调整，信号差→降权
    """
    
    SENSOR_WEIGHTS = {
        'hrv': 0.35,       # 最鲁棒，权重最高
        'breathing': 0.25, # 第二鲁棒
        'rppg': 0.15,      # 视觉非接触
        'eye_tracking': 0.10, # 光照敏感
        'eda': 0.05,       # 振动敏感，权重最低
        'behavior': 0.10   # 行为指标
    }
    
    def __init__(self):
        self.quality_scores = {k: 1.0 for k in self.SENSOR_WEIGHTS}
        
    def update_quality(self, sensor: str, quality: float):
        """更新传感器信号质量评分 (0-1)"""
        self.quality_scores[sensor] = max(0, min(1, quality))
        
    def get_weighted_fatigue_score(self, sensor_values: dict) -> float:
        """
        加权融合多传感器疲劳分数
        
        Returns:
            fatigue_score: 0(清醒) - 1(重度疲劳)
        """
        total_weight = 0
        weighted_sum = 0
        
        for sensor, value in sensor_values.items():
            w = self.SENSOR_WEIGHTS.get(sensor, 0)
            q = self.quality_scores.get(sensor, 0)
            effective_w = w * q
            weighted_sum += value * effective_w
            total_weight += effective_w
        
        if total_weight == 0:
            return 0.5  # 默认中等，不盲判断
            
        return weighted_sum / total_weight
    
    def report_sensor_status(self) -> dict:
        """报告各传感器状态"""
        return {
            sensor: {
                'quality': self.quality_scores[sensor],
                'weight': self.SENSOR_WEIGHTS[sensor],
                'effective_weight': self.SENSOR_WEIGHTS[sensor] * self.quality_scores[sensor],
                'status': 'active' if self.quality_scores[sensor] > 0.5 else 'degraded'
            }
            for sensor in self.SENSOR_WEIGHTS
        }