"""
goVISCAN多光谱rPPG原理

传统rPPG(远程光电容积描记):
- 用RGB摄像头拍摄面部
- 提取皮肤微小亮度变化(<1%)
- 亮度变化由血液脉动引起
- 绿光通道最敏感(血红蛋白吸收峰)

goVISCAN多光谱rPPG:
- 9个光谱通道同时成像
- 不同波长对血液不同成分敏感:
  * 绿光(~530nm): 血红蛋白吸收峰→心率
  * 红光(~620nm): 氧合/脱氧血红蛋白差异→SpO₂
  * NIR(~760-940nm): 穿透更深→呼吸(胸腔起伏)
- 多通道融合→更鲁棒、更准确
"""

import numpy as np
from typing import Tuple

class MultispectralrPPG:
    """多光谱rPPG信号提取"""
    
    def __init__(self, n_channels=9):
        self.n_channels = n_channels
        # 光谱权重（简化）
        # 通道0: 蓝(450nm), 1: 绿(530nm), 2: 红(620nm)
        # 通道3-5: NIR, 通道6-8: 扩展
        self.spectral_weights = np.array([
            0.3,  # 蓝 - 较弱rPPG信号
            1.0,  # 绿 - 最强rPPG(血红蛋白吸收峰)
            0.8,  # 红 - SpO₂计算
            0.6,  # NIR1 - 深层组织
            0.5,  # NIR2 - 呼吸
            0.4,  # NIR3 - 呼吸补充
            0.3, 0.2, 0.1  # 扩展
        ])
    
    def extract_heart_rate(self, multi_frames: np.ndarray, fps: int = 30) -> float:
        """
        从多光谱视频帧提取心率
        
        Args:
            multi_frames: (N, 9, H, W) N帧×9通道
            fps: 帧率
        
        Returns:
            heart_rate_bpm: 心率(次/分)
        """
        # 1. 皮肤ROI检测（简化：使用绿光通道）
        roi_mask = self._detect_skin(multi_frames[:, 1])
        
        # 2. 每通道提取时序信号
        signals = []
        for ch in range(self.n_channels):
            ch_frames = multi_frames[:, ch]
            mean_val = np.array([
                np.mean(ch_frames[i][roi_mask]) 
                for i in range(len(ch_frames))
            ])
            signals.append(mean_val)
        
        signals = np.array(signals)  # (9, N)
        
        # 3. 加权融合
        weighted = np.average(signals, axis=0, weights=self.spectral_weights)
        
        # 4. 去趋势
        weighted = weighted - np.convolve(weighted, np.ones(30)/30, mode='same')
        
        # 5. FFT找心率频率
        freqs = np.fft.rfftfreq(len(weighted), 1/fps)
        spectrum = np.abs(np.fft.rfft(weighted))
        
        # 心率范围: 0.7-3Hz (42-180 bpm)
        mask = (freqs > 0.7) & (freqs < 3.0)
        hr_freq = freqs[mask][np.argmax(spectrum[mask])]
        
        return hr_freq * 60  # Hz→bpm
    
    def extract_breathing(self, multi_frames: np.ndarray, fps: int = 30) -> float:
        """从NIR通道提取呼吸率"""
        # NIR通道对胸腔运动更敏感
        nir_channels = multi_frames[:, 3:6]  # 通道3-5
        
        # 提取低频信号(呼吸0.1-0.5Hz)
        signals = nir_channels.mean(axis=(2, 3))  # (N, 3)
        combined = signals.mean(axis=1)
        
        # 去趋势
        combined = combined - np.convolve(combined, np.ones(90)/90, mode='same')
        
        # FFT
        freqs = np.fft.rfftfreq(len(combined), 1/fps)
        spectrum = np.abs(np.fft.rfft(combined))
        
        mask = (freqs > 0.1) & (freqs < 0.5)
        br_freq = freqs[mask][np.argmax(spectrum[mask])]
        
        return br_freq * 60
    
    def extract_spo2(self, multi_frames: np.ndarray) -> float:
        """从红光+NIR计算血氧饱和度"""
        red_ch = multi_frames[:, 2].mean(axis=(1, 2))  # 红光
        nir_ch = multi_frames[:, 4].mean(axis=(1, 2))  # NIR
        
        # 比值法
        ac_red = red_ch - np.mean(red_ch)
        ac_nir = nir_ch - np.mean(nir_ch)
        
        ratio = np.std(ac_red) / max(np.std(ac_nir), 1e-6)
        spo2 = 110 - 25 * ratio  # 经验公式
        
        return max(70, min(100, spo2))
    
    def _detect_skin(self, frames):
        """简化皮肤检测"""
        return np.ones_like(frames[0], dtype=bool)


# 测试
if __name__ == "__main__":
    rppg = MultispectralrPPG()
    
    # 模拟数据
    np.random.seed(42)
    n_frames = 300  # 10秒@30fps
    frames = np.random.randn(n_frames, 9, 64, 64) * 10 + 128
    
    # 添加心率为72bpm的微弱脉动信号
    t = np.arange(n_frames) / 30
    pulse = np.sin(2 * np.pi * 1.2 * t) * 0.5  # 72bpm=1.2Hz
    for ch in [1, 2]:  # 绿光+红光
        frames[:, ch] += pulse[np.newaxis, np.newaxis, :] * 0.1
    
    hr = rppg.extract_heart_rate(frames)
    br = rppg.extract_breathing(frames)
    spo2 = rppg.extract_spo2(frames)
    
    print("=== goVISCAN多光谱rPPG测试 ===")
    print(f"心率: {hr:.1f} bpm (期望~72)")
    print(f"呼吸: {br:.1f} breaths/min")
    print(f"血氧: {spo2:.1f}%")
    print(f"\n论文报告精度(vs脉搏血氧计):")
    print(f"  心率误差: 2.8 bpm")
    print(f"  呼吸误差: 1.3 breaths/min")
    print(f"  血氧误差: 3.8 pp")