import numpy as np
from scipy import signal
from typing import Dict, List, Tuple, Optional

class MultimodalHeartbeatDetector:
    """
    车载多模态心率检测融合器
    
    基于 Nature Scientific Reports 2026 论文
    支持5种模态: ECG, PPG, iPPG-RGB, iPPG-IR, iPPG-Thermal
    
    融合策略:
    1. 各模态独立心跳检测
    2. 信号质量评估(SQA)
    3. 质量加权融合
    4. R-R间期一致性检查
    """
    
    def __init__(self, fs: int = 500):
        self.fs = fs
        self.modalities = {}
    
    def add_ecg_signal(self, ecg: np.ndarray):
        """方向盘ECG信号"""
        self.modalities['ecg'] = ecg
    
    def add_ppg_signal(self, ppg: np.ndarray):
        """方向盘PPG信号"""
        self.modalities['ppg'] = ppg
    
    def add_ippg_rgb(self, rgb_signal: np.ndarray):
        """RGB摄像头rPPG信号"""
        self.modalities['ippg_rgb'] = rgb_signal
    
    def add_ippg_ir(self, ir_signal: np.ndarray):
        """红外摄像头rPPG信号"""
        self.modalities['ippg_ir'] = ir_signal
    
    def add_ippg_thermal(self, thermal_signal: np.ndarray):
        """热成像rPPG信号"""
        self.modalities['ippg_thermal'] = thermal_signal
    
    def assess_quality(self, sig: np.ndarray, modality: str) -> float:
        """
        信号质量评估 (SQA)
        
        Returns:
            quality: 0-1, 越高越好
        """
        if len(sig) == 0:
            return 0.0
        
        # 1. 信噪比估计
        # 心率频段: 0.7-4.0 Hz (42-240 bpm)
        nyq = self.fs / 2
        low = 0.7 / nyq
        high = 4.0 / nyq
        
        # 带通滤波
        b, a = signal.butter(4, [low, high], btype='band')
        filtered = signal.filtfilt(b, a, sig)
        
        # 信号功率 vs 噪声功率
        signal_power = np.var(filtered)
        noise_power = np.var(sig - filtered)
        
        if noise_power < 1e-10:
            snr = 10.0
        else:
            snr = 10 * np.log10(signal_power / noise_power)
        
        # 2. 模态特定质量因子
        if modality == 'ecg':
            # ECG: 检查R波是否存在
            quality = min(1.0, max(0.0, (snr + 5) / 20))
        elif modality == 'ppg':
            # PPG: 检查脉动成分
            quality = min(1.0, max(0.0, (snr + 3) / 15))
        elif modality.startswith('ippg'):
            # rPPG: 更严格的质量要求
            quality = min(1.0, max(0.0, (snr + 2) / 12))
        else:
            quality = 0.0
        
        return quality
    
    def detect_heartbeats(self, sig: np.ndarray, 
                          modality: str) -> Tuple[np.ndarray, float]:
        """
        单模态心跳检测
        
        Returns:
            beat_times: R波/峰值时间点
            confidence: 检测置信度
        """
        if len(sig) == 0:
            return np.array([]), 0.0
        
        # 带通滤波
        nyq = self.fs / 2
        b, a = signal.butter(4, 
                            [0.7/nyq, 4.0/nyq], 
                            btype='band')
        filtered = signal.filtfilt(b, a, sig)
        
        # 模态特定的峰值检测
        if modality == 'ecg':
            # ECG: 使用Pan-Tompkins算法的简化版
            # 微分 + 平方 + 滑动窗口积分
            diff = np.diff(filtered)
            squared = diff ** 2
            window = int(0.15 * self.fs)
            integrated = np.convolve(squared, 
                                     np.ones(window)/window, 
                                     mode='same')
            peaks, _ = signal.find_peaks(integrated, 
                                          height=0.3*np.max(integrated),
                                          distance=0.3*self.fs)
        else:
            # PPG/rPPG: 直接峰值检测
            peaks, _ = signal.find_peaks(filtered,
                                          height=0.3*np.max(filtered),
                                          distance=0.3*self.fs)
        
        # 置信度
        quality = self.assess_quality(sig, modality)
        
        return peaks, quality
    
    def fuse(self, duration_sec: float = 10.0) -> Dict:
        """
        多模态融合心率检测
        
        Returns:
            result: 包含心率、HRV、各模态贡献的字典
        """
        all_beats = {}  # modality -> (beat_times, quality)
        
        # 各模态独立检测
        for modality, sig in self.modalities.items():
            beats, quality = self.detect_heartbeats(sig, modality)
            all_beats[modality] = (beats, quality)
        
        if not all_beats:
            return {'hr': 0, 'hrv': 0, 'modalities': {}}
        
        # 质量加权融合
        total_weight = sum(q for _, q in all_beats.values())
        if total_weight < 0.01:
            return {'hr': 0, 'hrv': 0, 'modalities': all_beats}
        
        # 加权平均心率
        weighted_hr = 0
        for modality, (beats, quality) in all_beats.items():
            if len(beats) > 1:
                rr_intervals = np.diff(beats) / self.fs
                hr = 60 / np.median(rr_intervals)
                weight = quality / total_weight
                weighted_hr += hr * weight
        
        # 融合HRV (使用最高质量模态的RR间期)
        best_modality = max(all_beats, 
                           key=lambda m: all_beats[m][1])
        best_beats, _ = all_beats[best_modality]
        if len(best_beats) > 2:
            rr = np.diff(best_beats) / self.fs
            hrv = np.std(rr) * 1000  # ms
        else:
            hrv = 0
        
        return {
            'hr': weighted_hr,
            'hrv': hrv,
            'best_modality': best_modality,
            'modalities': {m: {'beats': len(b), 'quality': q} 
                          for m, (b, q) in all_beats.items()}
        }


# 测试
if __name__ == "__main__":
    print("=" * 60)
    print("车载多模态心率检测融合器 - 测试")
    print("=" * 60)
    
    detector = MultimodalHeartbeatDetector(fs=500)
    
    # 模拟信号
    t = np.arange(0, 10, 1/500)
    heart_rate = 72  # bpm
    beat_interval = 60 / heart_rate  # seconds
    
    # ECG模拟 (R波)
    ecg = np.zeros_like(t)
    for i, beat_time in enumerate(np.arange(0, 10, beat_interval)):
        idx = int(beat_time * 500)
        if idx < len(ecg):
            ecg[max(0,idx-5):idx+5] = 1.0  # R波
    ecg += np.random.normal(0, 0.1, len(ecg))
    
    # PPG模拟
    ppg = 0.5 + 0.3 * np.sin(2 * np.pi * heart_rate/60 * t)
    ppg += np.random.normal(0, 0.05, len(ppg))
    
    # rPPG模拟 (弱信号)
    rppg = 0.1 * np.sin(2 * np.pi * heart_rate/60 * t)
    rppg += np.random.normal(0, 0.08, len(rppg))
    
    # 添加信号
    detector.add_ecg_signal(ecg)
    detector.add_ppg_signal(ppg)
    detector.add_ippg_rgb(rppg)
    detector.add_ippg_ir(rppg * 1.2)  # IR稍好
    detector.add_ippg_thermal(rppg * 0.8)  # 热成像稍差
    
    # 融合检测
    result = detector.fuse(duration_sec=10)
    
    print(f"\n融合结果:")
    print(f"  心率: {result['hr']:.1f} bpm (真实: {heart_rate} bpm)")
    print(f"  HRV: {result['hrv']:.1f} ms")
    print(f"  最优模态: {result['best_modality']}")
    print(f"\n各模态状态:")
    for m, info in result['modalities'].items():
        print(f"  {m:15s}: {info['beats']:3d} beats, quality={info['quality']:.3f}")