import numpy as np

def mimo_virtual_array(ntx: int, nrx: int) -> np.ndarray:
    """
    计算 MIMO 虚拟阵列位置
    
    Args:
        ntx: 发射天线数
        nrx: 接收天线数
    
    Returns:
        virtual_positions: 虚拟天线位置数组
    """
    # TX 间距为 d_tx，RX 间距为 d_rx
    # 虚拟阵列位置 = ntx * pos_rx + pos_tx
    d = 0.5  # 半波长间距 (λ/2 at 60GHz ≈ 2.5mm)
    
    tx_positions = np.arange(ntx) * d * nrx  # TX 间距放大 nrx 倍
    rx_positions = np.arange(nrx) * d
    
    virtual_positions = []
    for tx in tx_positions:
        for rx in rx_positions:
            virtual_positions.append(tx + rx)
    
    return np.array(sorted(virtual_positions))

# 对比
array_6t6r = mimo_virtual_array(6, 6)
array_3t4r = mimo_virtual_array(3, 4)
array_4t4r = mimo_virtual_array(4, 4)

print(f"6T6R 虚拟通道数: {len(array_6t6r)}")
print(f"3T4R 虚拟通道数: {len(array_3t4r)}")
print(f"4T4R 虚拟通道数: {len(array_4t4r)}")
print(f"6T6R vs 3T4R 角度分辨率提升: {len(array_6t6r)/len(array_3t4r):.1f}x")
print(f"6T6R vs 4T4R 角度分辨率提升: {len(array_6t6r)/len(array_4t4r):.1f}x")

# 输出:
# 6T6R 虚拟通道数: 36
# 3T4R 虚拟通道数: 12
# 4T4R 虚拟通道数: 16
# 6T6R vs 3T4R 角度分辨率提升: 3.0x
# 6T6R vs 4T4R 角度分辨率提升: 2.2x

import numpy as np
from scipy import signal as sig

def static_clutter_removal(range_profile: np.ndarray) -> np.ndarray:
    """
    静态杂波去除
    
    消除来自座椅、车身、天花板等静态反射体的信号，
    只保留动态成分（如呼吸引起的微小位移）。
    
    Args:
        range_profile: 距离向 FFT 结果, shape=(n_chirps, n_range_bins)
        
    Returns:
        dynamic_profile: 去除静态杂波后的动态距离剖面
    """
    # 方法1: 均值相消
    static_component = np.mean(range_profile, axis=0, keepdims=True)
    dynamic_profile = range_profile - static_component
    
    # 方法2: 高通滤波（可选，进一步提升）
    # 对每个距离门沿 chirp 维度做差分
    dynamic_profile = np.diff(dynamic_profile, axis=0, prepend=dynamic_profile[:1])
    
    return dynamic_profile


def extract_breathing_signal(
    phase_data: np.ndarray, 
    fps: int = 20,
    freq_min: float = 0.1,
    freq_max: float = 1.0
) -> np.ndarray:
    """
    从相位数据中提取呼吸信号
    
    Args:
        phase_data: 解卷绕后的相位序列, shape=(n_samples,)
        fps: 帧率（chirp 速率）
        freq_min: 呼吸频率下限 (Hz), 默认 0.1 (6次/分)
        freq_max: 呼吸频率上限 (Hz), 默认 1.0 (60次/分)
    
    Returns:
        breathing_signal: 滤波后的呼吸信号
    """
    nyquist = fps / 2
    b, a = sig.butter(4, [freq_min/nyquist, freq_max/nyquist], btype='band')
    breathing_signal = sig.filtfilt(b, a, phase_data)
    
    return breathing_signal


def detect_vital_signs(
    breathing_signal: np.ndarray,
    fps: int = 20,
    threshold: float = 0.5
) -> dict:
    """
    检测生命体征
    
    Args:
        breathing_signal: 滤波后的呼吸信号
        fps: 帧率
        threshold: 检测阈值（微伏级别位移）
    
    Returns:
        result: 检测结果
    """
    # FFT 分析频率成分
    freqs = np.fft.rfftfreq(len(breathing_signal), 1/fps)
    fft_mag = np.abs(np.fft.rfft(breathing_signal))
    
    # 在呼吸频带内寻找峰值
    mask = (freqs >= 0.1) & (freqs <= 1.0)
    peak_idx = np.argmax(fft_mag[mask])
    peak_freq = freqs[mask][peak_idx]
    peak_amp = fft_mag[mask][peak_idx]
    
    # 判断是否为活体
    is_alive = peak_amp > threshold
    
    return {
        'is_alive': is_alive,
        'breathing_rate': peak_freq * 60,  # 次/分钟
        'signal_amplitude': peak_amp,
        'confidence': min(peak_amp / (threshold * 2), 1.0)
    }


# 完整 CPD 检测流程测试
if __name__ == "__main__":
    np.random.seed(42)
    
    # 模拟 60GHz 雷达数据
    n_chirps = 200  # 10秒 @ 20fps
    n_range_bins = 64
    
    # 场景1: 座椅无人（只有静态杂波）
    range_data_empty = np.random.randn(n_chirps, n_range_bins) * 0.1 + 5.0
    
    # 场景2: 儿童在座椅上（有呼吸信号）
    t = np.linspace(0, 10, n_chirps)
    breathing = 0.3 * np.sin(2 * np.pi * 0.35 * t)  # 21次/分呼吸
    range_data_child = np.random.randn(n_chirps, n_range_bins) * 0.1 + 5.0
    range_data_child[:, 30] += breathing  # 距离门30处有呼吸信号
    
    # 场景3: 毛毯覆盖儿童（呼吸信号减弱）
    range_data_covered = range_data_child.copy()
    range_data_covered[:, 30] *= 0.3  # 信号衰减 70%
    
    for name, data in [("无人", range_data_empty), 
                        ("儿童", range_data_child),
                            ("毛毯覆盖", range_data_covered)]:
        dynamic = static_clutter_removal(data)
        phase = np.angle(np.fft.fft(dynamic, axis=0)[:, 30])
        phase = np.unwrap(phase)
        breathing = extract_breathing_signal(phase, fps=20)
        result = detect_vital_signs(breathing, fps=20, threshold=0.5)
        
        print(f"\n场景: {name}")
        print(f"  检测到活体: {result['is_alive']}")
        print(f"  呼吸频率: {result['breathing_rate']:.1f} 次/分")
        print(f"  信号幅度: {result['signal_amplitude']:.3f}")
        print(f"  置信度: {result['confidence']:.1%}")

# IMS CPD 模块集成伪代码（概念设计）

class CPDModule:
    """CPD 模块接口"""
    
    def __init__(self, radar_config: dict):
        self.radar = RadarSensor(
            chip="calterah_lancang_usrr",
            frequency=60e9,  # 60 GHz
            bandwidth=3.6e9,  # 3.6 GHz
            ntx=6, nrx=6,
            fps=20
        )
        self.clutter_filter = ClutterRemoval(method="mean_subtract")
        self.vital_detector = VitalSignDetector(
            freq_range=(0.1, 1.0),  # Hz
            min_amplitude=0.5,  # mm 微位移
            min_confidence=0.8
        )
        self.seat_map = SeatMapping()  # 距离门→座椅位置映射
    
    def run_cpd_cycle(self) -> dict:
        """执行一次 CPD 检测周期"""
        # 1. 获取雷达原始数据
        raw_data = self.radar.acquire(n_chirps=200)
        
        # 2. 信号处理
        range_profile = self.radar.range_fft(raw_data)
        dynamic = self.clutter_filter.remove(range_profile)
        phase = self.radar.extract_phase(dynamic)
        breathing = self.vital_detector.filter(phase)
        
        # 3. 生命体征检测
        result = self.vital_detector.detect(breathing)
        
        # 4. 座椅定位
        seat = self.seat_map.locate(result.range_bin)
        
        return {
            'detection': result.is_alive,
            'seat_position': seat,
            'breathing_rate': result.breathing_rate,
            'confidence': result.confidence,
            'timestamp': self.radar.timestamp
        }
    
    def get_cpd_status(self) -> dict:
        """获取 CPD 状态（供 IVI/HMI 显示）"""
        return {
            'status': 'monitoring',  # monitoring | alert | idle
            'rear_left': self._check_seat('rear_left'),
            'rear_right': self._check_seat('rear_right'),
            'rear_center': self._check_seat('rear_center'),
        }