"""
Fraunhofer goVISCAN 多光谱生命体征检测原理

核心原理：
1. 血液每次搏动引起皮肤微小颜色变化（人眼不可见）
2. 多光谱相机在蓝光到近红外多个波段同时采集
3. 9个微透镜各自对应不同光谱范围
4. 多通道分离提取心率、呼吸、血氧

对比传统 rPPG：
- 传统 rPPG: 单波长（绿光）→ 只能提取心率
- goVISCAN: 多波长 → 心率+呼吸+血氧
"""

import numpy as np
from typing import Tuple


class MultiSpectralVitalSigns:
    """
    多光谱生命体征提取器
    
    模拟 goVISCAN 的多通道 PPG 处理
    """
    
    # 9个微透镜对应的波长 (nm)
    WAVELENGTHS = [
        430,  # 蓝光 - 血红蛋白吸收峰
        470,  # 蓝光
        500,  # 青光
        530,  # 绿光 - rPPG 传统波长
        560,  # 绿光
        590,  # 黄光
        660,  # 红光 - SpO2 计算波长1
        730,  # 深红
        940,  # 近红外 - SpO2 计算波长2
    ]
    
    def __init__(self, num_wavelengths=9):
        self.num_wavelengths = num_wavelengths
        
    def extract_heart_rate(self, multi_spec_signal: np.ndarray,
                          fps: int = 30) -> float:
        """
        从多光谱信号提取心率
        
        多波长融合 PPG：
        1. 每个波长通道都有 PPG 信号
        2. 不同波长对不同肤色/光照有不同灵敏度
        3. 融合多波长提高鲁棒性
        
        Args:
            multi_spec_signal: shape=(num_wavelengths, time_steps)
            fps: 采样率
            
        Returns:
            heart_rate: 心率 (次/分)
        """
        # 归一化每个通道
        normalized = np.zeros_like(multi_spec_signal)
        for i in range(self.num_wavelengths):
            channel = multi_spec_signal[i]
            normalized[i] = (channel - np.mean(channel)) / \
                            (np.std(channel) + 1e-8)
        
        # 加权融合（绿光波段权重最高）
        weights = np.array([
            0.3, 0.5, 0.7, 1.0, 1.0, 0.8, 0.6, 0.5, 0.4
        ])
        fused = np.average(normalized, axis=0, weights=weights)
        
        # 带通滤波 0.8-3.0 Hz (48-180 bpm)
        from scipy.signal import butter, filtfilt
        b, a = butter(4, [0.8, 3.0], btype='band', fs=fps)
        filtered = filtfilt(b, a, fused)
        
        # FFT 找主频
        freqs = np.fft.rfftfreq(len(filtered), 1/fps)
        spectrum = np.abs(np.fft.rfft(filtered))
        dominant_idx = np.argmax(spectrum[1:]) + 1
        heart_rate = freqs[dominant_idx] * 60
        
        return heart_rate
    
    def extract_respiratory_rate(self, multi_spec_signal: np.ndarray,
                                 fps: int = 30) -> float:
        """
        从多光谱信号提取呼吸频率
        
        呼吸引起的皮肤色调变化频率更低（0.1-0.5 Hz）
        """
        # 低频成分提取
        from scipy.signal import butter, filtfilt
        b, a = butter(4, [0.1, 0.5], btype='band', fs=fps)
        
        # 融合
        fused = np.mean(multi_spec_signal, axis=0)
        filtered = filtfilt(b, a, fused)
        
        freqs = np.fft.rfftfreq(len(filtered), 1/fps)
        spectrum = np.abs(np.fft.rfft(filtered))
        dominant_idx = np.argmax(spectrum[1:]) + 1
        resp_rate = freqs[dominant_idx] * 60
        
        return resp_rate
    
    def extract_spo2(self, multi_spec_signal: np.ndarray) -> float:
        """
        从多光谱信号计算血氧饱和度
        
        使用红光 (660nm) 和近红外 (940nm) 的比值
        SpO2 = f(R/IR ratio)
        
        标准公式：
        SpO2 = 110 - 25 * (AC_red/DC_red) / (AC_ir/DC_ir)
        """
        # 红光通道 (660nm, index 6)
        red_channel = multi_spec_signal[6]
        # 近红外通道 (940nm, index 8)
        ir_channel = multi_spec_signal[8]
        
        # AC 分量（脉动）
        red_ac = np.std(red_channel)
        ir_ac = np.std(ir_channel)
        
        # DC 分量（基线）
        red_dc = np.mean(red_channel)
        ir_dc = np.mean(ir_channel)
        
        # 比值
        ratio = (red_ac / red_dc) / (ir_ac / ir_dc + 1e-8)
        
        # 经验公式
        spo2 = 110 - 25 * ratio
        
        return np.clip(spo2, 70, 100)
    
    def anti_motion_artifact(self, multi_spec_signal: np.ndarray) -> np.ndarray:
        """
        运动伪影消除
        
        goVISCAN 的核心优势：
        9个微透镜同时采集，可区分运动伪影和真实PPG
        真实PPG在所有波长上同步变化
        运动伪影在不同波长上变化不同
        """
        # 计算各通道间的相关性
        corr_matrix = np.corrcoef(multi_spec_signal)
        
        # 使用 PCA 提取主要成分（真实PPG）
        # 运动伪影在次要成分中
        from numpy.linalg import svd
        U, S, Vt = svd(multi_spec_signal, full_matrices=False)
        
        # 第一主成分通常是真实PPG
        clean_signal = U[:, 0:1] @ np.diag(S[0:1]) @ Vt[0:1, :]
        
        # 添加次要成分的部分（保留部分信息）
        clean_signal += 0.3 * U[:, 1:2] @ np.diag(S[1:2]) @ Vt[1:2, :]
        
        return clean_signal


# ========== 模拟测试 ==========

if __name__ == "__main__":
    np.random.seed(42)
    extractor = MultiSpectralVitalSigns()
    
    # 模拟多光谱信号
    duration = 30  # 30秒
    fps = 30
    num_frames = duration * fps
    t = np.linspace(0, duration, num_frames)
    
    # 真实参数
    true_hr = 72  # 心率 72 bpm = 1.2 Hz
    true_rr = 16  # 呼吸 16 bpm = 0.267 Hz
    true_spo2 = 97  # 血氧 97%
    
    # 生成9通道信号
    multi_signal = np.zeros((9, num_frames))
    for i in range(9):
        # PPG 信号（各波长幅度不同）
        ppg = np.sin(2*np.pi*true_hr/60 * t) * (0.5 + 0.5*np.random.rand())
        # 呼吸调制
        resp = np.sin(2*np.pi*true_rr/60 * t) * 0.2
        # 运动伪影
        motion = 0.3 * np.sin(2*np.pi*0.7*t + np.random.rand()*2*np.pi)
        # 噪声
        noise = 0.1 * np.random.randn(num_frames)
        
        multi_signal[i] = ppg + resp + motion + noise
    
    # 提取生命体征
    hr = extractor.extract_heart_rate(multi_signal, fps)
    rr = extractor.extract_respiratory_rate(multi_signal, fps)
    spo2 = extractor.extract_spo2(multi_signal)
    
    print("=== goVISCAN 多光谱生命体征检测 ===")
    print(f"  真实心率: {true_hr} bpm, 检测: {hr:.1f} bpm, "
          f"误差: {abs(hr-true_hr):.1f}")
    print(f"  真实呼吸: {true_rr} bpm, 检测: {rr:.1f} bpm, "
          f"误差: {abs(rr-true_rr):.1f}")
    print(f"  真实血氧: {true_spo2}%, 检测: {spo2:.1f}%")
    
    # 对比 goVISCAN 实验室精度
    print(f"\n=== goVISCAN 实验室精度（论文报告）===")
    print(f"  心率误差: 2.8 beats/min")
    print(f"  呼吸误差: 1.3 breaths/min")
    print(f"  血氧误差: 3.8 percentage points")