import numpy as np
import matplotlib.pyplot as plt

class FMCWRadar:
    """FMCW 毫米波雷达仿真
    
    参数（参考 TI IWR6843AOP）：
    - 中心频率：60.5 GHz
    - 带宽：4 GHz
    - 扫频周期：100 μs
    - 最大探测距离：5 m
    """
    
    def __init__(self):
        # 雷达参数
        self.fc = 60.5e9       # 中心频率 60.5 GHz
        self.B = 4e9           # 带宽 4 GHz
        self.Tc = 100e-6       # 扫频周期 100 μs
        self.c = 3e8           # 光速
        self.fs = 10e6         # 采样率 10 MHz
        
        # 计算参数
        self.max_range = 5.0   # 最大距离 5m
        self.range_resolution = self.c / (2 * self.B)  # 距离分辨率 ~3.75cm
        
    def generate_chirp(self):
        """生成线性调频信号
        
        Returns:
            tx_signal: 发射信号
            time: 时间轴
        """
        num_samples = int(self.Tc * self.fs)
        time = np.linspace(0, self.Tc, num_samples)
        
        # 线性调频：频率从 fc 变化到 fc+B
        # φ(t) = 2π * (fc*t + 0.5*K*t^2), K = B/Tc
        K = self.B / self.Tc  # 调频斜率
        phase = 2 * np.pi * (self.fc * time + 0.5 * K * time**2)
        
        tx_signal = np.exp(1j * phase)
        
        return tx_signal, time
    
    def simulate_target_reflection(self, tx_signal, distance, velocity):
        """仿真目标反射
        
        Args:
            tx_signal: 发射信号
            distance: 目标距离 (米)
            velocity: 目标速度 (米/秒)
            
        Returns:
            rx_signal: 接收信号
            beat_freq: 拍频 (Hz)
        """
        # 时延
        tau = 2 * distance / self.c
        
        # 多普勒频移
        fd = 2 * velocity * self.fc / self.c
        
        # 接收信号（延迟 + 多普勒）
        # 简化：只考虑相位延迟
        delay_samples = int(tau * self.fs)
        rx_signal = np.zeros_like(tx_signal)
        rx_signal[delay_samples:] = tx_signal[:-delay_samples] * \
                                      np.exp(1j * 2 * np.pi * fd * 
                                             np.arange(len(tx_signal) - delay_samples) / self.fs)
        
        # 拍频（中频信号）
        K = self.B / self.Tc
        beat_freq = K * tau  # 距离对应的拍频
        
        return rx_signal, beat_freq
    
    def range_fft(self, rx_signal):
        """距离 FFT（Range FFT）
        
        将时域信号转换为距离域，提取目标距离信息
        
        Returns:
            range_profile: 距离剖面（幅度谱）
            range_axis: 距离轴 (米)
        """
        # FFT
        fft_result = np.fft.fft(rx_signal)
        fft_magnitude = np.abs(fft_result)
        
        # 距离轴
        num_samples = len(fft_result)
        range_axis = np.linspace(0, self.max_range, num_samples)
        
        return fft_magnitude, range_axis
    
    def detect_vital_signs(self, rx_signal, num_frames=100):
        """生命体征检测
        
        通过连续测量相位变化，提取呼吸和心跳频率
        
        Args:
            rx_signal: 接收信号序列（多帧）
            num_frames: 帧数
            
        Returns:
            respiration_freq: 呼吸频率 (Hz)
            heart_rate_freq: 心跳频率 (Hz)
        """
        # 简化实现：假设已知相位序列
        # 实际需要完整信号处理流程
        
        # 模拟呼吸（0.2-0.5 Hz，12-30 次/分）
        respiration_freq = 0.3  # Hz
        
        # 模拟心跳（1.0-1.5 Hz，60-90 次/分）
        heart_rate_freq = 1.2  # Hz
        
        return respiration_freq, heart_rate_freq


# 实际测试
if __name__ == "__main__":
    radar = FMCWRadar()
    
    # 生成发射信号
    tx_signal, time = radar.generate_chirp()
    
    # 仿真目标反射（距离 1.5m，胸壁运动模拟呼吸）
    distance = 1.5  # 米
    velocity = 0.01  # 呼吸时胸壁运动速度 ~1 cm/s
    
    rx_signal, beat_freq = radar.simulate_target_reflection(tx_signal, distance, velocity)
    
    # 距离 FFT
    range_profile, range_axis = radar.range_fft(rx_signal)
    
    # 输出
    print(f"距离分辨率: {radar.range_resolution * 100:.2f} cm")
    print(f"拍频: {beat_freq / 1e6:.2f} MHz")
    print(f"检测距离: {distance:.2f} m")
    print(f"呼吸频率: {radar.detect_vital_signs(rx_signal)[0]:.2f} Hz")
    print(f"心跳频率: {radar.detect_vital_signs(rx_signal)[1]:.2f} Hz")

class VitalSignsExtractor:
    """生命体征提取器
    
    从雷达相位信号中提取呼吸和心跳频率
    
    参考：
    - Nature Scientific Reports, 2025: "Detection of vital signs based on millimeter wave radar"
    - MDPI Electronics, 2024: "Noncontact Monitoring of Respiration and Heartbeat"
    """
    
    def __init__(self, fs=20):
        """
        Args:
            fs: 采样率 (Hz)，典型值 20 Hz
        """
        self.fs = fs
        
        # 频率范围
        self.respiration_band = (0.1, 0.5)   # Hz (6-30 次/分)
        self.heart_rate_band = (0.8, 2.0)    # Hz (48-120 次/分)
        
    def extract_phase(self, radar_frames):
        """从雷达数据中提取相位信息
        
        Args:
            radar_frames: 雷达帧序列, shape=(num_frames, num_samples)
            
        Returns:
            phase_signal: 相位信号, shape=(num_frames,)
        """
        # 对每帧做距离 FFT
        num_frames, num_samples = radar_frames.shape
        
        phase_signal = np.zeros(num_frames)
        
        for i in range(num_frames):
            # 距离 FFT
            range_fft = np.fft.fft(radar_frames[i])
            
            # 找到目标距离对应的峰值（假设在索引 50）
            target_idx = 50
            
            # 提取相位
            phase_signal[i] = np.angle(range_fft[target_idx])
        
        # 相位展开（unwrap）
        phase_signal = np.unwrap(phase_signal)
        
        return phase_signal
    
    def remove_clutter(self, phase_signal):
        """去除静态杂波和呼吸谐波
        
        Args:
            phase_signal: 相位信号
            
        Returns:
            cleaned_signal: 清理后的信号
        """
        from scipy import signal
        
        # 1. 去除直流分量（静态杂波）
        phase_signal = phase_signal - np.mean(phase_signal)
        
        # 2. 带通滤波（保留呼吸和心跳频段）
        nyquist = self.fs / 2
        low_cut = 0.1  # Hz
        high_cut = 2.0  # Hz
        
        sos = signal.butter(4, [low_cut / nyquist, high_cut / nyquist],
                           btype='band', output='sos')
        cleaned_signal = signal.sosfiltfilt(sos, phase_signal)
        
        return cleaned_signal
    
    def estimate_respiration(self, phase_signal):
        """估计呼吸频率
        
        Args:
            phase_signal: 相位信号
            
        Returns:
            respiration_freq: 呼吸频率 (Hz)
            respiration_rate: 呼吸频率 (次/分)
        """
        # FFT
        n = len(phase_signal)
        freq = np.fft.fftfreq(n, 1/self.fs)
        fft_magnitude = np.abs(np.fft.fft(phase_signal))
        
        # 只取正频率
        pos_freq = freq[:n//2]
        pos_magnitude = fft_magnitude[:n//2]
        
        # 在呼吸频段内找峰值
        respiration_idx = np.where((pos_freq >= self.respiration_band[0]) &
                                    (pos_freq <= self.respiration_band[1]))[0]
        
        peak_idx = respiration_idx[np.argmax(pos_magnitude[respiration_idx])]
        respiration_freq = pos_freq[peak_idx]
        respiration_rate = respiration_freq * 60  # 次/分
        
        return respiration_freq, respiration_rate
    
    def estimate_heart_rate(self, phase_signal, respiration_freq):
        """估计心跳频率
        
        需要先去除呼吸谐波干扰
        
        Args:
            phase_signal: 相位信号
            respiration_freq: 已估计的呼吸频率
            
        Returns:
            heart_rate_freq: 心跳频率 (Hz)
            heart_rate_bpm: 心跳频率 (次/分)
        """
        # 创建自适应滤波器去除呼吸谐波
        from scipy import signal
        
        # 简化：直接在心跳频段搜索
        n = len(phase_signal)
        freq = np.fft.fftfreq(n, 1/self.fs)
        fft_magnitude = np.abs(np.fft.fft(phase_signal))
        
        pos_freq = freq[:n//2]
        pos_magnitude = fft_magnitude[:n//2]
        
        # 在心跳频段内找峰值（排除呼吸谐波）
        heart_rate_idx = np.where((pos_freq >= self.heart_rate_band[0]) &
                                   (pos_freq <= self.heart_rate_band[1]))[0]
        
        peak_idx = heart_rate_idx[np.argmax(pos_magnitude[heart_rate_idx])]
        heart_rate_freq = pos_freq[peak_idx]
        heart_rate_bpm = heart_rate_freq * 60
        
        return heart_rate_freq, heart_rate_bpm
    
    def process(self, radar_frames):
        """完整处理流程
        
        Args:
            radar_frames: 雷达帧序列
            
        Returns:
            vital_signs: dict
                - respiration_rate: 呼吸频率 (次/分)
                - heart_rate: 心跳频率 (次/分)
                - confidence: 置信度
        """
        # 提取相位
        phase_signal = self.extract_phase(radar_frames)
        
        # 去除杂波
        cleaned_signal = self.remove_clutter(phase_signal)
        
        # 估计呼吸
        resp_freq, resp_rate = self.estimate_respiration(cleaned_signal)
        
        # 估计心跳
        hr_freq, hr_bpm = self.estimate_heart_rate(cleaned_signal, resp_freq)
        
        return {
            'respiration_rate': resp_rate,
            'heart_rate': hr_bpm,
            'confidence': self._calculate_confidence(cleaned_signal)
        }
    
    def _calculate_confidence(self, signal):
        """计算检测置信度"""
        # 简化：基于信号能量
        energy = np.sum(signal**2)
        confidence = min(energy / 100, 1.0)  # 归一化
        return confidence


# 实际测试
if __name__ == "__main__":
    import numpy as np
    
    # 模拟雷达数据（100 帧）
    np.random.seed(42)
    radar_frames = np.random.randn(100, 256) * 0.1
    
    # 处理
    extractor = VitalSignsExtractor(fs=20)
    vital_signs = extractor.process(radar_frames)
    
    print(f"呼吸频率: {vital_signs['respiration_rate']:.1f} 次/分")
    print(f"心跳频率: {vital_signs['heart_rate']:.1f} 次/分")
    print(f"置信度: {vital_signs['confidence']:.2f}")

class IMSCPDModule:
    """IMS CPD 检测模块（雷达+视觉融合）"""
    
    def __init__(self):
        # 雷达处理器
        self.radar_processor = VitalSignsExtractor()
        
        # 视觉检测器（预训练模型）
        self.vision_detector = None  # 实际需要加载模型
        
        # 融合权重
        self.radar_weight = 0.6
        self.vision_weight = 0.4
        
    def detect_child_presence(self, radar_data, ir_frame):
        """检测儿童存在
        
        Args:
            radar_data: 雷达数据, shape=(num_frames, num_samples)
            ir_frame: 红外图像, shape=(H, W)
            
        Returns:
            result: dict
                - is_child_present: bool
                - vital_signs: dict
                - confidence: float
        """
        # 1. 雷达生命体征检测
        vital_signs = self.radar_processor.process(radar_data)
        
        # 2. 视觉儿童检测（简化）
        vision_confidence = self._vision_detect_child(ir_frame)
        
        # 3. 融合决策
        # 关键：呼吸频率在儿童范围（30-60 次/分）且视觉检测到小物体
        is_child = (vital_signs['respiration_rate'] > 30 and
                    vital_signs['respiration_rate'] < 60 and
                    vital_signs['confidence'] > 0.7)
        
        # 4. 综合置信度
        final_confidence = (vital_signs['confidence'] * self.radar_weight +
                           vision_confidence * self.vision_weight)
        
        return {
            'is_child_present': is_child,
            'vital_signs': vital_signs,
            'confidence': final_confidence
        }
    
    def _vision_detect_child(self, ir_frame):
        """视觉儿童检测（简化实现）"""
        # 实际需要：
        # 1. 目标检测（儿童/婴儿座椅）
        # 2. 尺度估计
        # 3. 姿态估计（躺/坐）
        return 0.8  # 示例值