import numpy as np
from dataclasses import dataclass
from typing import Optional, Tuple
from scipy.signal import butter, filtfilt, find_peaks
from scipy.stats import zscore

"""
车载健康监测三路融合系统
方向盘ECG + 安全带呼吸 + 摄像头rPPG

参考德国研究团队方案，将在 Automechanika Frankfurt 2026 展出
"""

@dataclass
class HealthMetrics:
    """健康指标"""
    # 心率（多源）
    hr_ecg: float         # 方向盘ECG心率
    hr_rppg: float        # 摄像头rPPG心率
    hr_fused: float       # 融合心率
    
    # HRV
    sdnn: float           # RR间期标准差
    rmssd: float          # 连续差值均方根
    lf_hf: float          # 交感/副交感平衡
    
    # 呼吸
    resp_rate: float      # 呼吸频率(次/分)
    resp_pattern: str     # 呼吸模式
    
    # 综合
    stress_index: float   # 压力指数 0-1
    health_score: float   # 健康评分 0-100
    anomaly_flag: bool    # 异常标志
    anomaly_type: str     # 异常类型


class SteeringWheelECG:
    """方向盘 ECG 传感器模拟"""
    
    def __init__(self, sample_rate: int = 500):
        self.fs = sample_rate
        
    def process(self, ecg_signal: np.ndarray) -> dict:
        """
        处理 ECG 信号
        
        Args:
            ecg_signal: 方向盘采集的ECG信号
            
        Returns:
            {'hr': float, 'rr_intervals': np.ndarray, 'sdnn': float, ...}
        """
        # 带通滤波 0.5-40 Hz
        nyq = self.fs / 2
        b, a = butter(3, [0.5/nyq, 40/nyq], btype='band')
        ecg_filtered = filtfilt(b, a, ecg_signal)
        
        # QRS 检测（简化版）
        # 放大R峰
        diff = np.diff(ecg_filtered)
        squared = diff ** 2
        # 滑动平均
        window = int(self.fs * 0.15)
        smoothed = np.convolve(squared, np.ones(window)/window, mode='same')
        
        # 峰值检测
        peaks, _ = find_peaks(smoothed, distance=self.fs*0.4, 
                               height=np.percentile(smoothed, 75))
        
        if len(peaks) < 3:
            return {'hr': 0, 'rr_intervals': np.array([]),
                    'sdnn': 0, 'rmssd': 0, 'quality': 0}
        
        rr = np.diff(peaks) / self.fs * 1000  # ms
        
        # HRV 特征
        hr = 60000 / np.mean(rr)
        sdnn = np.std(rr)
        diff_rr = np.diff(rr)
        rmssd = np.sqrt(np.mean(diff_rr**2)) if len(diff_rr) > 0 else 0
        
        # 信号质量
        quality = min(1.0, len(peaks) / (len(ecg_signal) / self.fs * 1.2))
        
        # 心律失常检测（简化）
        arrhythmia = False
        if len(rr) > 5:
            rr_z = np.abs(zscore(rr))
            if np.any(rr_z > 2.5):
                arrhythmia = True
        
        return {
            'hr': round(hr, 1),
            'rr_intervals': rr,
            'sdnn': round(sdnn, 1),
            'rmssd': round(rmssd, 1),
            'quality': round(quality, 3),
            'arrhythmia': arrhythmia
        }


class SeatbeltRespiratory:
    """安全带呼吸传感器模拟"""
    
    def __init__(self, sample_rate: int = 50):
        self.fs = sample_rate
    
    def process(self, tension_signal: np.ndarray) -> dict:
        """
        处理安全带张力信号
        
        Args:
            tension_signal: 安全带张力变化信号
            
        Returns:
            {'resp_rate': float, 'pattern': str, 'quality': float}
        """
        # 带通 0.1-0.7 Hz (6-42 breaths/min)
        nyq = self.fs / 2
        b, a = butter(2, [0.1/nyq, 0.7/nyq], btype='band')
        resp_filtered = filtfilt(b, a, tension_signal)
        
        # 峰值检测
        peaks, _ = find_peaks(resp_filtered, distance=self.fs*1.5,
                               height=np.std(resp_filtered))
        
        if len(peaks) < 3:
            return {'resp_rate': 0, 'pattern': 'unknown', 'quality': 0}
        
        # 呼吸频率
        rr_intervals = np.diff(peaks) / self.fs  # 秒
        resp_rate = 60 / np.mean(rr_intervals)
        
        # 模式判断
        rr_std = np.std(rr_intervals)
        if resp_rate < 12:
            pattern = 'bradypnea'  # 呼吸过缓
        elif resp_rate > 20:
            pattern = 'tachypnea'  # 呼吸过速
        elif rr_std > 1.5:
            pattern = 'irregular'
        else:
            pattern = 'normal'
        
        quality = min(1.0, len(peaks) / (len(tension_signal) / self.fs * 0.2))
        
        return {
            'resp_rate': round(resp_rate, 1),
            'pattern': pattern,
            'quality': round(quality, 3),
            'amplitude': round(np.std(resp_filtered), 3)
        }


class DMSrPPG:
    """DMS 摄像头 rPPG 模拟"""
    
    def __init__(self, fps: int = 30):
        self.fps = fps
    
    def process(self, face_frames: np.ndarray) -> dict:
        """从面部帧提取 rPPG"""
        # 简化版 POS 算法
        rgb_means = np.array([
            f.reshape(-1, 3).mean(axis=0) for f in face_frames
        ])
        
        S = 3 * rgb_means[:, 0] - 2 * rgb_means[:, 1]
        H = 1.5 * rgb_means[:, 0] - rgb_means[:, 1]
        
        S = (S - np.mean(S)) / (np.std(S) + 1e-8)
        H = (H - np.mean(H)) / (np.std(H) + 1e-8)
        
        ppg = S + 0.3 * H
        
        # 带通
        nyq = self.fps / 2
        b, a = butter(3, [0.7/nyq, 3.5/nyq], btype='band')
        ppg_filtered = filtfilt(b, a, ppg)
        
        peaks, _ = find_peaks(ppg_filtered, distance=self.fps*0.4)
        
        if len(peaks) < 3:
            return {'hr': 0, 'sdnn': 0, 'quality': 0}
        
        rr = np.diff(peaks) / self.fps * 1000
        hr = 60000 / np.mean(rr)
        sdnn = np.std(rr)
        quality = min(1.0, len(peaks) / (len(face_frames) / self.fps * 1.5))
        
        return {
            'hr': round(hr, 1),
            'sdnn': round(sdnn, 1),
            'quality': round(quality, 3)
        }


class HealthFusionEngine:
    """三路健康数据融合引擎"""
    
    def __init__(self):
        self.ecg_sensor = SteeringWheelECG()
        self.seatbelt_sensor = SeatbeltRespiratory()
        self.rppg_sensor = DMSrPPG()
        self.baseline = None
        
    def fuse(self, ecg_data: dict, resp_data: dict, 
             rppg_data: dict, vehicle_data: dict = None) -> HealthMetrics:
        """三路融合"""
        
        # 1. 心率融合（加权平均，基于质量）
        hr_ecg = ecg_data.get('hr', 0)
        hr_rppg = rppg_data.get('hr', 0)
        q_ecg = ecg_data.get('quality', 0)
        q_rppg = rppg_data.get('quality', 0)
        
        if q_ecg > 0.5 and q_rppg > 0.5:
            # 两者都可靠，加权融合
            w_ecg = q_ecg / (q_ecg + q_rppg)
            hr_fused = hr_ecg * w_ecg + hr_rppg * (1 - w_ecg)
        elif q_ecg > 0.5:
            hr_fused = hr_ecg
        elif q_rppg > 0.5:
            hr_fused = hr_rppg
        else:
            hr_fused = 0
        
        # 2. HRV（优先用 ECG）
        sdnn = ecg_data.get('sdnn', rppg_data.get('sdnn', 0))
        rmssd = ecg_data.get('rmssd', 0)
        
        # 3. 呼吸
        resp_rate = resp_data.get('resp_rate', 0)
        resp_pattern = resp_data.get('pattern', 'unknown')
        
        # 4. 压力指数
        stress = self._compute_stress(hr_fused, sdnn, resp_rate)
        
        # 5. 健康评分
        score = self._compute_health_score(hr_fused, sdnn, resp_rate, 
                                           ecg_data.get('arrhythmia', False))
        
        # 6. 异常检测
        anomaly, anomaly_type = self._detect_anomaly(
            hr_fused, resp_rate, resp_pattern,
            ecg_data.get('arrhythmia', False), stress
        )
        
        return HealthMetrics(
            hr_ecg=hr_ecg, hr_rppg=hr_rppg, hr_fused=round(hr_fused, 1),
            sdnn=round(sdnn, 1), rmssd=round(rmssd, 1),
            lf_hf=round(sdnn / (rmssd + 1e-8), 2),
            resp_rate=resp_rate, resp_pattern=resp_pattern,
            stress_index=round(stress, 3),
            health_score=round(score, 1),
            anomaly_flag=anomaly, anomaly_type=anomaly_type
        )
    
    def _compute_stress(self, hr, sdnn, resp_rate):
        """综合压力指数"""
        hr_comp = max(0, min(1, (hr - 60) / 40))
        hrv_comp = max(0, min(1, 1 - sdnn / 50))
        resp_comp = max(0, min(1, (resp_rate - 12) / 12))
        return hr_comp * 0.4 + hrv_comp * 0.4 + resp_comp * 0.2
    
    def _compute_health_score(self, hr, sdnn, resp_rate, arrhythmia):
        """健康评分 0-100"""
        if arrhythmia:
            return 40.0
        score = 100
        if hr < 50 or hr > 100:
            score -= 20
        if sdnn < 20:
            score -= 15
        if resp_rate < 10 or resp_rate > 25:
            score -= 15
        return max(0, score)
    
    def _detect_anomaly(self, hr, resp_rate, resp_pattern, 
                        arrhythmia, stress):
        """异常检测"""
        if arrhythmia:
            return True, 'cardiac_arrhythmia'
        if hr > 120 or hr < 40:
            return True, 'heart_rate_abnormal'
        if resp_rate > 25 or resp_rate < 8:
            return True, 'respiratory_abnormal'
        if stress > 0.8:
            return True, 'stress_overload'
        return False, 'none'


# ==================== 测试 ====================
if __name__ == "__main__":
    np.random.seed(42)
    
    print("=" * 70)
    print("车载健康监测三路融合系统测试")
    print("方向盘ECG + 安全带呼吸 + DMS rPPG")
    print("=" * 70)
    
    engine = HealthFusionEngine()
    
    # 模拟正常驾驶
    fs_ecg = 500
    fs_resp = 50
    fps = 30
    duration = 10  # 秒
    
    # 正常 ECG 信号 (70 bpm)
    t_ecg = np.arange(fs_ecg * duration) / fs_ecg
    ecg_normal = np.sin(2 * np.pi * (70/60) * t_ecg) * 0.5 + np.random.normal(0, 0.05, len(t_ecg))
    
    # 正常呼吸 (16 breaths/min = 0.27 Hz)
    t_resp = np.arange(fs_resp * duration) / fs_resp
    resp_normal = np.sin(2 * np.pi * 0.27 * t_resp) * 0.3 + np.random.normal(0, 0.02, len(t_resp))
    
    # 正常 rPPG (70 bpm)
    t_rppg = np.arange(fps * duration) / fps
    frames_normal = np.random.randint(100, 120, (fps * duration, 50, 50, 3), dtype=np.uint8)
    pulse = np.sin(2 * np.pi * (70/60) * t_rppg) * 2
    for i in range(len(frames_normal)):
        frames_normal[i] = np.clip(frames_normal[i].astype(float) + pulse[i], 0, 255).astype(np.uint8)
    
    # 处理
    ecg_result = engine.ecg_sensor.process(ecg_normal)
    resp_result = engine.seatbelt_sensor.process(resp_normal)
    rppg_result = engine.rppg_sensor.process(frames_normal)
    
    print("\n--- 正常驾驶 ---")
    print(f"  方向盘ECG: HR={ecg_result['hr']} bpm, SDNN={ecg_result['sdnn']}ms, 质量={ecg_result['quality']}")
    print(f"  安全带呼吸: {resp_result['resp_rate']} 次/分, 模式={resp_result['pattern']}, 质量={resp_result['quality']}")
    print(f"  rPPG: HR={rppg_result['hr']} bpm, 质量={rppg_result['quality']}")
    
    metrics = engine.fuse(ecg_result, resp_result, rppg_result)
    print(f"\n  融合结果:")
    print(f"    心率: ECG={metrics.hr_ecg} | rPPG={metrics.hr_rppg} | 融合={metrics.hr_fused}")
    print(f"    HRV: SDNN={metrics.sdnn}ms | RMSSD={metrics.rmssd}ms | LF/HF={metrics.lf_hf}")
    print(f"    呼吸: {metrics.resp_rate} 次/分 ({metrics.resp_pattern})")
    print(f"    压力指数: {metrics.stress_index}")
    print(f"    健康评分: {metrics.health_score}/100")
    print(f"    异常: {metrics.anomaly_flag} ({metrics.anomaly_type})")
    
    # 模拟压力状态
    ecg_stress = np.sin(2 * np.pi * (95/60) * t_ecg) * 0.5 + np.random.normal(0, 0.08, len(t_ecg))
    resp_stress = np.sin(2 * np.pi * 0.38 * t_resp) * 0.2 + np.random.normal(0, 0.03, len(t_resp))  # 23 breaths/min
    
    ecg_s = engine.ecg_sensor.process(ecg_stress)
    resp_s = engine.seatbelt_sensor.process(resp_stress)
    
    metrics_s = engine.fuse(ecg_s, resp_s, rppg_result)
    
    print(f"\n--- 压力状态 ---")
    print(f"  ECG: HR={ecg_s['hr']} bpm, SDNN={ecg_s['sdnn']}ms")
    print(f"  呼吸: {resp_s['resp_rate']} 次/分 ({resp_s['pattern']})")
    print(f"  融合心率: {metrics_s.hr_fused} bpm")
    print(f"  压力指数: {metrics_s.stress_index}")
    print(f"  健康评分: {metrics_s.health_score}/100")
    print(f"  异常: {metrics_s.anomaly_flag} ({metrics_s.anomaly_type})")
    
    print(f"\n{'='*70}")
    print("传感器覆盖矩阵:")
    print(f"  {'状态':<15} {'方向盘ECG':>12} {'安全带':>12} {'rPPG':>12}")
    print(f"  {'双手握盘':<15} {'✅':>12} {'✅':>12} {'✅':>12}")
    print(f"  {'单手驾驶':<15} {'❌':>12} {'✅':>12} {'✅':>12}")
    print(f"  {'脱手(ADAS)':<15} {'❌':>12} {'✅':>12} {'✅':>12}")
    print(f"  {'夜间':<15} {'✅':>12} {'✅':>12} {'⚠️需IR':>12}")
    print(f"  {'戴手套':<15} {'❌':>12} {'✅':>12} {'✅':>12}")
    print(f"{'='*70}")