"""
多站UWB同步与数据融合仿真
基于IEEE 802.15.4ab空口同步协议

硬件配置：
- Calterah Alps系列 UWB芯片
- 2T4R 时分MIMO（2发射 + 4接收）
- 同步精度要求：< 1ns（皮秒级时间精度）
- 工作频段：6-8 GHz UWB
"""

import numpy as np
from dataclasses import dataclass
from typing import List, Tuple

@dataclass
class UW BAnchor:
    """UWB锚节点定义"""
    anchor_id: int
    position: np.ndarray  # 3D位置 (x, y, z) in meters
    is_transmitter: bool
    is_receiver: bool
    sync_offset_ns: float = 0.0  # 同步偏移

@dataclass
class CPDTarget:
    """CPD检测目标"""
    position: np.ndarray
    breathing_rate: float  # breaths per minute
    chest_displacement_mm: float  # 胸部位移幅度 mm
    is_child: bool

class MultistaticUWBCPD:
    """
    多站UWB儿童存在检测系统
    
    基于Calterah Alps芯片 + IEEE 802.15.4ab同步
    """
    
    def __init__(self, anchors: List[UW BAnchor]):
        self.anchors = anchors
        self.sync_precision_ns = 0.5  # 同步精度 0.5ns
        self.frequency = 7.5e9  # 7.5 GHz 中心频率
        self.bandwidth = 500e6  # 500 MHz 带宽
        
    def compute_bistatic_range(
        self, tx_pos: np.ndarray, rx_pos: np.ndarray, target_pos: np.ndarray
    ) -> float:
        """
        计算双站距离（发射-目标-接收路径长度）
        
        Args:
            tx_pos: 发射节点位置
            rx_pos: 接收节点位置
            target_pos: 目标位置
            
        Returns:
            bistatic_range: 双站距离 (m)
        """
        path_tx = np.linalg.norm(target_pos - tx_pos)
        path_rx = np.linalg.norm(rx_pos - target_pos)
        return path_tx + path_rx
    
    def localize_target(
        self, 
        measurements: List[Tuple[int, int, float]]
    ) -> np.ndarray:
        """
        多站三角定位
        
        Args:
            measurements: [(tx_id, rx_id, bistatic_range), ...]
            
        Returns:
            position: 目标3D位置估计
        """
        from scipy.optimize import least_squares
        
        def residuals(pos):
            res = []
            for tx_id, rx_id, measured_range in measurements:
                tx = self.anchors[tx_id]
                rx = self.anchors[rx_id]
                predicted = self.compute_bistatic_range(
                    tx.position, rx.position, pos
                )
                res.append(predicted - measured_range)
            return np.array(res)
        
        # 初始估计：车内中心
        x0 = np.array([1.5, 0.0, 0.5])
        result = least_squares(residuals, x0)
        return result.x
    
    def detect_breathing(
        self, 
        phase_signal: np.ndarray, 
        fs: float = 24e6
    ) -> dict:
        """
        从UWB相位信号中提取呼吸频率
        
        基于微多普勒分析：
        - 成人呼吸: 12-20 bpm, 胸部位移 4-12mm
        - 婴儿呼吸: 30-60 bpm, 胸部位移 2-5mm
        - 心跳: 60-150 bpm, 胸部位移 0.2-0.5mm
        
        Args:
            phase_signal: 相位信号序列
            fs: 采样率
            
        Returns:
            dict: 呼吸频率、心率、置信度
        """
        from scipy.signal import find_peaks, butter, filtfilt
        
        # 带通滤波：0.1-2 Hz (覆盖呼吸+心跳)
        nyq = fs / 2
        b, a = butter(4, [0.1/nyq, 2.0/nyq], btype='band')
        filtered = filtfilt(b, a, phase_signal)
        
        # FFT提取频率成分
        fft = np.fft.rfft(filtered)
        freqs = np.fft.rfftfreq(len(filtered), 1/fs)
        magnitude = np.abs(fft)
        
        # 呼吸频段 (0.2-1.0 Hz = 12-60 bpm)
        breath_mask = (freqs >= 0.2) & (freqs <= 1.0)
        breath_freq = freqs[breath_mask][np.argmax(magnitude[breath_mask])]
        breath_rate_bpm = breath_freq * 60
        
        # 心跳频段 (1.0-2.5 Hz = 60-150 bpm)
        heart_mask = (freqs >= 1.0) & (freqs <= 2.5)
        if np.any(magnitude[heart_mask] > 0.1 * np.max(magnitude)):
            heart_freq = freqs[heart_mask][np.argmax(magnitude[heart_mask])]
            heart_rate_bpm = heart_freq * 60
        else:
            heart_rate_bpm = None
        
        # 判断是否为儿童
        is_child = breath_rate_bpm > 25  # 婴儿/儿童呼吸频率 > 25 bpm
        
        # 置信度评估
        snr = np.max(magnitude[breath_mask]) / np.mean(magnitude)
        confidence = min(1.0, snr / 10.0)
        
        return {
            'breathing_rate_bpm': breath_rate_bpm,
            'heart_rate_bpm': heart_rate_bpm,
            'is_child': is_child,
            'confidence': confidence,
            'snr_db': 10 * np.log10(snr)
        }


# ==================== 实际测试 ====================
if __name__ == "__main__":
    # 车内UWB锚节点布局（复用数字钥匙硬件）
    # 典型位置：前中控、左B柱、右B柱、后中控
    anchors = [
        UW BAnchor(0, np.array([2.0, 0.0, 1.2]), True, True),   # 前中控
        UW BAnchor(1, np.array([0.5, -0.8, 1.0]), False, True),  # 左B柱
        UW BAnchor(2, np.array([0.5, 0.8, 1.0]), False, True),  # 右B柱
        UW BAnchor(3, np.array([0.0, 0.0, 1.1]), True, False),  # 后中控
    ]
    
    cpd = MultistaticUWBCPD(anchors)
    
    # 模拟儿童位置（后排左侧座椅）
    child = CPDTarget(
        position=np.array([0.8, -0.5, 0.3]),
        breathing_rate=35,  # 婴儿呼吸 35 bpm
        chest_displacement_mm=3.0,
        is_child=True
    )
    
    # 模拟多站测量
    measurements = []
    for i, tx in enumerate(anchors):
        if not tx.is_transmitter:
            continue
        for j, rx in enumerate(anchors):
            if not rx.is_receiver or i == j:
                continue
            # 添加测量噪声
            true_range = cpd.compute_bistatic_range(
                tx.position, rx.position, child.position
            )
            measured = true_range + np.random.normal(0, 0.02)  # 2cm噪声
            measurements.append((i, j, measured))
    
    # 定位
    estimated_pos = cpd.localize_target(measurements)
    error = np.linalg.norm(estimated_pos - child.position)
    print(f"真实位置: {child.position}")
    print(f"估计位置: {estimated_pos}")
    print(f"定位误差: {error*100:.1f} cm")
    
    # 模拟呼吸信号
    t = np.linspace(0, 10, 2400)  # 10秒, 240Hz
    breath_signal = 3.0 * np.sin(2 * np.pi * 35/60 * t)  # 35 bpm
    noise = np.random.normal(0, 0.5, len(t))
    phase_signal = breath_signal + noise
    
    result = cpd.detect_breathing(phase_signal, fs=240)
    print(f"\n--- 生命体征检测结果 ---")
    print(f"呼吸频率: {result['breathing_rate_bpm']:.1f} bpm")
    print(f"心率: {result['heart_rate_bpm']}")
    print(f"判断为儿童: {result['is_child']}")
    print(f"置信度: {result['confidence']:.2f}")
    print(f"信噪比: {result['snr_db']:.1f} dB")