import numpy as np
from scipy.signal import butter, filtfilt

class DopplerRespiratoryAmplifier:
    """
    多普勒域呼吸放大器
    
    论文：arXiv 2026-01
    用于FMCW雷达半静态人体检测
    """
    
    def __init__(self, fs=20, f_low=0.2, f_high=0.5):
        """
        Args:
            fs: 多普勒采样率 (Hz)
            f_low: 呼吸频段下限 (Hz)
            f_high: 呼吸频段上限 (Hz)
        """
        self.fs = fs
        # 带通滤波器（呼吸频段）
        self.b, self.a = butter(4, [f_low/(fs/2), f_high/(fs/2)], btype='band')
    
    def amplify_respiration(self, doppler_signal):
        """
        放大呼吸信号
        
        Args:
            doppler_signal: (N_samples,) 多普勒域原始信号
        
        Returns:
            amplified: 放大后的呼吸信号
            energy: 呼吸频段能量
        """
        # 带通滤波提取呼吸频段
        filtered = filtfilt(self.b, self.a, doppler_signal)
        
        # 能量累积
        energy = np.sum(filtered ** 2) / len(filtered)
        
        # 放大（归一化到0-1）
        if energy > 0:
            amplified = filtered / np.max(np.abs(filtered))
        else:
            amplified = filtered
        
        return amplified, energy
    
    def detect_occupant(self, doppler_signal, threshold=1e-6):
        """
        检测半静态乘员
        
        Args:
            doppler_signal: (N_samples,) 多普勒信号
            threshold: 检测阈值
        
        Returns:
            is_occupied: bool
            energy: 呼吸频段能量
        """
        _, energy = self.amplify_respiration(doppler_signal)
        is_occupied = energy > threshold
        return is_occupied, energy


# 测试
if __name__ == "__main__":
    detector = DopplerRespiratoryAmplifier(fs=20)
    
    # 模拟空座椅信号（噪声）
    np.random.seed(42)
    noise = np.random.randn(200) * 0.001
    occupied, energy = detector.detect_occupant(noise)
    print(f"空座椅: 检测={'有人' if occupied else '空'}, 能量={energy:.2e}")
    
    # 模拟呼吸信号（0.3Hz + 噪声）
    t = np.arange(200) / 20
    breathing = 0.01 * np.sin(2 * np.pi * 0.3 * t) + np.random.randn(200) * 0.001
    occupied, energy = detector.detect_occupant(breathing)
    print(f"呼吸人体: 检测={'有人' if occupied else '空'}, 能量={energy:.2e}")