"""
BCG (Ballistocardiography) 座舱心率监测系统
从座椅传感器原始信号到心率/HRV输出的完整管道

参考论文：
1. Nature Scientific Data (2025): BCG dataset with ECG reference
2. MDPI Sensors 26(2):548 (2026): MM-TFNet BCG+PPG fusion
3. Annals of Biomedical Engineering (2026): Critical review

关键指标：
- 心率精度：±2-5 BPM（vs ECG参考）
- HRV精度：IBI误差 < 30ms
- 心搏量偏差：1.58 mL（最佳研究, Direct Fick参考）
- 心输出量偏差：0.31 L/min
- 信号带宽：1-15 Hz
- 信噪比：> 10 dB（静止坐姿）
"""

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

@dataclass
class BCGConfig:
    """BCG 系统配置"""
    # 传感器参数
    sampling_rate: float = 1000.0  # Hz
    sensor_sensitivity: float = 0.1  # mV/mN
    
    # 信号处理参数
    bcg_band_low: float = 1.0   # Hz
    bcg_band_high: float = 15.0  # Hz
    notch_freq: float = 50.0  # 工频陷波 Hz
    
    # 心率提取参数
    hr_window_sec: int = 10  # 心率计算窗口
    hr_update_rate: int = 1  # 每秒更新
    min_hr: float = 40.0  # BPM
    max_hr: float = 200.0  # BPM
    
    # HRV参数
    hrv_window_sec: int = 60  # HRV计算窗口（需要更长）
    
    # 平均参数
    ensemble_beats: int = 8  # 集平均心跳数


class BCGProcessor:
    """
    BCG 信号处理器
    
    完整管道：
    1. 预处理（去基线、工频陷波、带通滤波）
    2. R-peak 等效检测（BCG J-peak）
    3. 心率计算
    4. HRV 参数提取
    5. 质量评估
    
    座舱特殊挑战：
    - 运动伪影（车辆振动、身体移动）
    - 呼吸干扰（0.15-0.5 Hz 与 BCG 重叠）
    - 多乘员干扰（前排+后排传感器串扰）
    - 温度漂移
    """
    
    def __init__(self, config: BCGConfig):
        self.config = config
        
    def preprocess(self, raw_signal: np.ndarray) -> np.ndarray:
        """
        BCG 预处理管道
        
        1. 去基线漂移（0.5 Hz 高通）
        2. 工频陷波（50/60 Hz）
        3. BCG 带通滤波（1-15 Hz）
        
        关键：使用线性相位滤波器避免时延失真！
        参考论文要求："linear-phase filters to avoid distorting 
        the timing of fiducial points"
        """
        fs = self.config.sampling_rate
        
        # 1. 去基线（高通 0.5 Hz）
        b_hp, a_hp = butter(4, 0.5 / (fs/2), btype='high')
        baseline_removed = filtfilt(b_hp, a_hp, raw_signal)
        
        # 2. 工频陷波
        from scipy.signal import iirnotch
        Q = 30  # 品质因子
        b_notch, a_notch = iirnotch(
            self.config.notch_freq / (fs/2), Q
        )
        notched = filtfilt(b_notch, a_notch, baseline_removed)
        
        # 3. BCG 带通（1-15 Hz，线性相位=零相位 filtfilt）
        b_bp, a_bp = butter(
            4, 
            [self.config.bcg_band_low / (fs/2),
             self.config.bcg_band_high / (fs/2)],
            btype='band'
        )
        bcg_filtered = filtfilt(b_bp, a_bp, notched)
        
        return bcg_filtered
    
    def detect_j_peaks(self, bcg_signal: np.ndarray,
                        ecg_r_peaks: Optional[np.ndarray] = None
                        ) -> np.ndarray:
        """
        检测 BCG J-peak（等效于 ECG R-peak）
        
        BCG 波形特征：
        - I峰：心室收缩前
        - J峰：主动脉血流峰值（~R-peak后 80-120ms）
        - K峰：减速期
        
        J-peak 是最可靠的检测点
        
        Args:
            bcg_signal: 预处理后的BCG信号
            ecg_r_peaks: 可选的ECG R-peak（用于验证/校准）
        
        Returns:
            j_peaks: J-peak 位置索引数组
        """
        # 自适应阈值
        signal_envelope = np.abs(bcg_signal)
        window = int(2.0 * self.config.sampling_rate)  # 2秒窗口
        adaptive_threshold = np.convolve(
            signal_envelope,
            np.ones(window) / window,
            mode='same'
        ) * 0.6
        
        # 峰检测
        min_distance = int(0.3 * self.config.sampling_rate)  # 最小心跳间隔 300ms
        peaks, properties = find_peaks(
            bcg_signal,
            height=adaptive_threshold,
            distance=min_distance,
            prominence=adaptive_threshold * 0.5
        )
        
        # 如果有ECG参考，计算 BCG J-peak 与 ECG R-peak 的时延
        if ecg_r_peaks is not None and len(peaks) > 0:
            delays = []
            for r_peak in ecg_r_peaks:
                # 找最近的 J-peak
                j_candidates = peaks[peaks > r_peak]
                if len(j_candidates) > 0:
                    j_peak = j_candidates[0]
                    delay_ms = (j_peak - r_peak) / self.config.sampling_rate * 1000
                    delays.append(delay_ms)
            
            if delays:
                mean_delay = np.mean(delays)
                print(f"  BCG J-peak 延迟: {mean_delay:.1f} ms (正常: 80-120ms)")
        
        return peaks
    
    def compute_heart_rate(self, j_peaks: np.ndarray) -> float:
        """
        从 J-peak 间隔计算心率
        
        Args:
            j_peaks: J-peak 索引数组
        
        Returns:
            heart_rate: BPM
        """
        if len(j_peaks) < 2:
            return 0.0
        
        # 计算间隔（秒）
        intervals = np.diff(j_peaks) / self.config.sampling_rate
        
        # 过滤异常间隔
        valid_mask = (
            (intervals > 60.0 / self.config.max_hr) &  # 最小间隔
            (intervals < 60.0 / self.config.min_hr)    # 最大间隔
        )
        valid_intervals = intervals[valid_mask]
        
        if len(valid_intervals) < 2:
            return 0.0
        
        # 平均间隔 → 心率
        mean_interval = np.mean(valid_intervals)
        heart_rate = 60.0 / mean_interval
        
        return heart_rate
    
    def compute_hrv(self, j_peaks: np.ndarray) -> dict:
        """
        计算 HRV 参数
        
        从 J-peak 间隔提取心率变异性
        
        返回：
        - SDNN: 间隔标准差（ms）
        - RMSSD: 连续间隔差值均方根（ms）
        - pNN50: 间隔差 >50ms 的比例
        - LF_HF_ratio: 低频/高频功率比
        """
        if len(j_peaks) < 10:
            return {'error': 'insufficient_peaks'}
        
        # IBI 间隔（毫秒）
        ibi = np.diff(j_peaks) / self.config.sampling_rate * 1000
        
        # 过滤异常
        valid_mask = (ibi > 300) & (ibi < 1500)
        ibi = ibi[valid_mask]
        
        if len(ibi) < 5:
            return {'error': 'insufficient_valid_ibi'}
        
        # 时域 HRV
        sdnn = np.std(ibi)
        rmssd = np.sqrt(np.mean(np.diff(ibi)**2))
        nn50 = np.sum(np.abs(np.diff(ibi)) > 50)
        pnn50 = nn50 / len(ibi) * 100
        
        # 频域 HRV（需要更长序列）
        if len(ibi) >= 30:
            # 插值为均匀采样
            t = np.cumsum(ibi) / 1000.0  # 秒
            t_uniform = np.arange(t[0], t[-1], 1.0/4.0)  # 4Hz
            ibi_uniform = np.interp(t_uniform, t, ibi)
            
            # 频谱分析
            freqs = np.fft.rfftfreq(len(ibi_uniform), 1.0/4.0)
            spectrum = np.abs(np.fft.rfft(ibi_uniform - np.mean(ibi_uniform)))
            
            # 功率频段
            lf_mask = (freqs >= 0.04) & (freqs < 0.15)
            hf_mask = (freqs >= 0.15) & (freqs < 0.40)
            
            lf_power = np.sum(spectrum[lf_mask]**2)
            hf_power = np.sum(spectrum[hf_mask]**2)
            lf_hf_ratio = lf_power / (hf_power + 1e-8)
        else:
            lf_power = hf_power = lf_hf_ratio = 0
        
        return {
            'mean_ibi_ms': np.mean(ibi),
            'sdnn_ms': sdnn,
            'rmssd_ms': rmssd,
            'pnn50_pct': pnn50,
            'lf_power': lf_power,
            'hf_power': hf_power,
            'lf_hf_ratio': lf_hf_ratio,
            'num_beats': len(ibi),
        }
    
    def ensemble_average(self, bcg_signal: np.ndarray,
                         j_peaks: np.ndarray,
                         pre_ms: int = 200,
                         post_ms: int = 400) -> np.ndarray:
        """
        集平均：对齐多个心跳取平均，抑制噪声
        
        参考论文："average eight to sixteen consecutive 
        R-wave-triggered beats to suppress respiratory noise"
        
        Args:
            bcg_signal: BCG信号
            j_peaks: J-peak位置
            pre_ms: J-peak前毫秒数
            post_ms: J-peak后毫秒数
        
        Returns:
            ensemble: 集平均波形
        """
        fs = self.config.sampling_rate
        pre_samples = int(pre_ms / 1000 * fs)
        post_samples = int(post_ms / 1000 * fs)
        beat_length = pre_samples + post_samples
        
        beats = []
        for peak in j_peaks:
            start = peak - pre_samples
            end = peak + post_samples
            if start >= 0 and end < len(bcg_signal):
                beat = bcg_signal[start:end]
                # 归一化（去直流）
                beat = beat - np.mean(beat)
                beats.append(beat)
        
        if len(beats) < self.config.ensemble_beats:
            return np.zeros(beat_length)
        
        # 集平均
        ensemble = np.mean(beats, axis=0)
        
        return ensemble
    
    def assess_signal_quality(self, bcg_signal: np.ndarray,
                               j_peaks: np.ndarray) -> dict:
        """
        评估BCG信号质量
        
        质量影响因素：
        1. 运动伪影（最大干扰源，行走时 > 10x 心信号）
        2. 呼吸干扰（0.15-0.5 Hz）
        3. 传感器位置偏移（2-3cm 偏移即显著影响波形）
        4. 多乘员串扰
        5. 温度漂移
        """
        # SNR估算
        if len(j_peaks) < 2:
            return {'quality': 'poor', 'snr_db': 0, 'confidence': 0}
        
        # 信号功率（心跳窗口内）
        fs = self.config.sampling_rate
        beat_windows = []
        for peak in j_peaks:
            start = max(0, peak - int(0.1 * fs))
            end = min(len(bcg_signal), peak + int(0.1 * fs))
            beat_windows.append(bcg_signal[start:end])
        
        signal_power = np.mean([np.var(w) for w in beat_windows])
        
        # 噪声功率（心跳间间隔）
        noise_windows = []
        for i in range(len(j_peaks) - 1):
            start = j_peaks[i] + int(0.2 * fs)
            end = j_peaks[i+1] - int(0.2 * fs)
            if end > start and end < len(bcg_signal):
                noise_windows.append(bcg_signal[start:end])
        
        if noise_windows:
            noise_power = np.mean([np.var(w) for w in noise_windows])
        else:
            noise_power = signal_power * 0.1
        
        snr = 10 * np.log10(signal_power / (noise_power + 1e-8))
        
        # 质量分级
        if snr > 15:
            quality = 'excellent'
            confidence = 0.95
        elif snr > 10:
            quality = 'good'
            confidence = 0.80
        elif snr > 5:
            quality = 'fair'
            confidence = 0.60
        else:
            quality = 'poor'
            confidence = 0.30
        
        # 心跳间隔一致性
        intervals = np.diff(j_peaks) / fs
        interval_cv = np.std(intervals) / (np.mean(intervals) + 1e-8)
        
        return {
            'quality': quality,
            'snr_db': snr,
            'confidence': confidence,
            'interval_cv': interval_cv,
            'num_beats': len(j_peaks),
        }


# ==================== 完整测试 ====================
if __name__ == "__main__":
    np.random.seed(42)
    
    config = BCGConfig()
    processor = BCGProcessor(config)
    
    print(f"=== BCG 座舱心率监测系统 ===")
    print(f"采样率: {config.sampling_rate} Hz")
    print(f"信号频段: {config.bcg_band_low}-{config.bcg_band_high} Hz")
    print(f"心率范围: {config.min_hr}-{config.max_hr} BPM")
    
    # 模拟 60秒 BCG 信号
    duration_sec = 60
    n_samples = int(config.sampling_rate * duration_sec)
    t = np.arange(n_samples) / config.sampling_rate
    
    # 真实心率 72 BPM = 1.2 Hz
    true_hr = 72.0
    true_ibi = 60.0 / true_hr  # 0.833 sec
    
    # 生成 BCG 波形（J-peak 为主峰）
    bcg_clean = np.zeros(n_samples)
    beat_interval_samples = int(true_ibi * config.sampling_rate)
    
    for i, peak_idx in enumerate(range(0, n_samples, beat_interval_samples)):
        # BCG J-peak 波形（简化高斯）
        for j in range(-50, 100):
            idx = peak_idx + j
            if 0 <= idx < n_samples:
                # IJK 复合波
                if -50 < j <= 0:  # I 波
                    bcg_clean[idx] -= 0.3 * np.exp(-(j+25)**2 / 50)
                elif 0 < j <= 30:  # J 波（主峰）
                    bcg_clean[idx] += 1.0 * np.exp(-j**2 / 100)
                elif 30 < j <= 80:  # K 波
                    bcg_clean[idx] -= 0.4 * np.exp(-(j-55)**2 / 200)
    
    # 加入噪声
    # 1. 呼吸干扰（0.3 Hz）
    breathing = 0.2 * np.sin(2 * np.pi * 0.3 * t)
    # 2. 车辆振动（多频段）
    vehicle_vibration = 0.1 * np.sin(2*np.pi*5*t) + 0.05*np.sin(2*np.pi*8*t)
    # 3. 白噪声
    white_noise = np.random.normal(0, 0.1, n_samples)
    # 4. 基线漂移
    baseline = 0.5 * np.sin(2 * np.pi * 0.05 * t)
    
    raw_signal = bcg_clean + breathing + vehicle_vibration + white_noise + baseline
    
    print(f"\n=== 信号生成 ===")
    print(f"时长: {duration_sec} 秒")
    print(f"真实心率: {true_hr} BPM")
    print(f"心跳数: ~{int(duration_sec * true_hr / 60)}")
    print(f"SNR (加入噪声后): {10*np.log10(np.var(bcg_clean)/np.var(breathing+vehicle_vibration+white_noise)):.1f} dB")
    
    # 预处理
    filtered = processor.preprocess(raw_signal)
    
    # J-peak 检测
    j_peaks = processor.detect_j_peaks(filtered)
    print(f"\n=== J-peak 检测 ===")
    print(f"检测到 J-peak 数: {len(j_peaks)}")
    print(f"预期心跳数: {int(duration_sec * true_hr / 60)}")
    
    # 心率
    hr = processor.compute_heart_rate(j_peaks)
    print(f"\n=== 心率 ===")
    print(f"估计心率: {hr:.1f} BPM")
    print(f"真实心率: {true_hr:.1f} BPM")
    print(f"误差: {abs(hr - true_hr):.1f} BPM")
    
    # HRV
    hrv = processor.compute_hrv(j_peaks)
    print(f"\n=== HRV ===")
    print(f"平均 IBI: {hrv.get('mean_ibi_ms', 0):.1f} ms (真实: {true_ibi*1000:.1f} ms)")
    print(f"SDNN: {hrv.get('sdnn_ms', 0):.1f} ms")
    print(f"RMSSD: {hrv.get('rmssd_ms', 0):.1f} ms")
    print(f"pNN50: {hrv.get('pnn50_pct', 0):.1f}%")
    print(f"LF/HF比: {hrv.get('lf_hf_ratio', 0):.2f}")
    
    # 信号质量
    quality = processor.assess_signal_quality(filtered, j_peaks)
    print(f"\n=== 信号质量 ===")
    print(f"质量等级: {quality['quality']}")
    print(f"SNR: {quality['snr_db']:.1f} dB")
    print(f"置信度: {quality['confidence']:.2f}")
    print(f"间隔变异系数: {quality['interval_cv']:.3f}")
    
    # 集平均
    ensemble = processor.ensemble_average(filtered, j_peaks)
    print(f"\n=== 集平均 ===")
    print(f"平均波形长度: {len(ensemble)} samples")
    print(f"使用心跳数: {min(len(j_peaks), config.ensemble_beats)}")
    
    print(f"\n=== 临床验证状态（来自综述） ===")
    print(f"心搏量偏差: 1.58 mL (Direct Fick 参考)")
    print(f"心输出量偏差: 0.31 L/min")
    print(f"趋势一致性: 96.7%")
    print(f"⚠️ 射血分数: 无 Bland-Altman 验证")
    print(f"⚠️ 血压: 未达 ISO 81060-2 标准")
    print(f"⚠️ 所有研究均为单中心、横断面")

"""
MM-TFNet: BCG + PPG 多模态时序融合网络
参考: MDPI Sensors 26(2):548 (2026)

核心方法：
- BCG 提供心搏量/心输出量信息
- PPG 提供脉搏波传导时间
- 融合后提升心率和HRV精度
"""

class MM_TFNet:
    """
    多模态时序融合网络（简化版）
    
    架构：
    1. BCG 分支：1D-CNN 提取振动特征
    2. PPG 分支：1D-CNN 提取光学特征  
    3. 时序融合：Cross-attention 对齐
    4. 输出头：心率 + HRV + 质量评分
    """
    
    def __init__(self):
        # BCG 分支参数
        self.bcg_filters = [8, 16, 32, 64]
        self.bcg_kernel = 7
        
        # PPG 分支参数
        self.ppg_filters = [8, 16, 32, 64]
        self.ppg_kernel = 5
        
        # 融合参数
        self.fusion_dim = 128
        self.num_heads = 4
        
    def forward_concept(self, bcg_window: np.ndarray,
                         ppg_window: np.ndarray) -> dict:
        """
        概念性前向传播（实际使用PyTorch）
        
        Args:
            bcg_window: (1, 1000) 1秒BCG
            ppg_window: (1, 1000) 1秒PPG
        
        Returns:
            output: {heart_rate, hrv_sdnn, quality_score}
        """
        # BCG 特征提取
        bcg_features = self._extract_temporal_features(
            bcg_window, self.bcg_filters, self.bcg_kernel
        )
        
        # PPG 特征提取
        ppg_features = self._extract_temporal_features(
            ppg_window, self.ppg_filters, self.ppg_kernel
        )
        
        # 跨模态注意力
        fused = self._cross_attention(
            bcg_features, ppg_features
        )
        
        # 输出
        return {
            'heart_rate': 72.0,  # BPM
            'hrv_sdnn': 45.0,  # ms
            'quality_score': 0.85,  # 0-1
            'fusion_confidence': 0.90,
        }
    
    def _extract_temporal_features(self, signal, filters, kernel):
        """1D-CNN 特征提取（简化）"""
        return np.random.randn(128)  # 概念输出
    
    def _cross_attention(self, feat_a, feat_b):
        """跨模态注意力融合"""
        return np.random.randn(256)