import numpy as np
from typing import Tuple

class UWBRadarProcessor:
    """
    Calterah UWB 雷达信号处理器（概念实现）
    
    基于分层参数设计：
    - 默认参数：适配大部分车型/环境
    - 微调参数：少量参数针对特定车型
    """
    
    def __init__(self, config: str = 'default'):
        """
        Args:
            config: 'default' 或 'custom_vehicle'
        """
        # 默认参数（适用大部分车型）
        self.default_params = {
            'range_resolution': 0.015,  # 1.5cm 距离分辨率
            'velocity_resolution': 0.05,  # 5cm/s 速度分辨率
            'detection_range': 2.0,  # 2m 检测范围
            'clutter_removal': 'mean_subtract',
            'static_threshold': 0.3,
            'breathing_freq_range': (0.3, 0.8),  # 18-48 次/分钟
            'heart_rate_range': (1.0, 2.0),  # 60-120 次/分钟
        }
        
        # 微调参数（仅少量需要按车型调整）
        self.tuning_params = {
            'antenna_delay': 0,  # 天线延迟（按安装位置）
            'clutter_map': None,  # 杂波图（按车内环境）
        }
    
    def process_cir(self, cir_data: np.ndarray) -> Dict:
        """
        处理 Channel Impulse Response 数据
        
        Args:
            cir_data: CIR 数据, shape=(N_frames, N_rx, N_samples)
            
        Returns:
            特征字典
        """
        # 1. 杂波去除
        clean_data = self._remove_clutter(cir_data)
        
        # 2. 距离-多普勒图
        range_doppler = self._compute_range_doppler(clean_data)
        
        # 3. 提取生命体征
        breathing = self._extract_breathing(clean_data)
        heart_rate = self._extract_heart_rate(clean_data)
        
        # 4. 特征向量
        features = {
            'range_profile': self._range_profile(clean_data),
            'doppler_profile': range_doppler,
            'breathing_rate': breathing,
            'breathing_amplitude': np.max(np.abs(breathing)),
            'heart_rate': heart_rate,
            'motion_level': np.var(clean_data, axis=0).mean(),
            'static_energy': np.mean(np.abs(clean_data)),
        }
        
        return features
    
    def _remove_clutter(self, data: np.ndarray) -> np.ndarray:
        """均值减法去杂波"""
        return data - np.mean(data, axis=0, keepdims=True)
    
    def _compute_range_doppler(self, data: np.ndarray) -> np.ndarray:
        """计算距离-多普勒图"""
        # 沿时间轴 FFT 得到多普勒
        doppler = np.fft.fft(data, axis=0)
        # 沿距离轴 FFT
        range_doppler = np.fft.fft(doppler, axis=-1)
        return np.abs(range_doppler)
    
    def _extract_breathing(self, data: np.ndarray) -> np.ndarray:
        """提取呼吸信号（带通滤波 0.3-0.8 Hz）"""
        from scipy.signal import butter, filtfilt
        
        fs = 20  # 假设 20fps
        low, high = self.default_params['breathing_freq_range']
        b, a = butter(4, [low/(fs/2), high/(fs/2)], btype='band')
        
        # 对每个接收通道的每个距离门滤波
        breathing = np.zeros_like(data)
        for rx in range(data.shape[1]):
            for rng in range(data.shape[2]):
                breathing[:, rx, rng] = filtfilt(b, a, data[:, rx, rng])
        
        return breathing
    
    def _extract_heart_rate(self, data: np.ndarray) -> np.ndarray:
        """提取心率信号（带通滤波 1.0-2.0 Hz）"""
        from scipy.signal import butter, filtfilt
        
        fs = 20
        low, high = self.default_params['heart_rate_range']
        b, a = butter(4, [low/(fs/2), high/(fs/2)], btype='band')
        
        hr = np.zeros_like(data)
        for rx in range(data.shape[1]):
            for rng in range(data.shape[2]):
                hr[:, rx, rng] = filtfilt(b, a, data[:, rx, rng])
        
        return hr
    
    def _range_profile(self, data: np.ndarray) -> np.ndarray:
        """距离剖面"""
        return np.abs(np.fft.fft(data, axis=-1)).mean(axis=0)


# 测试
if __name__ == "__main__":
    np.random.seed(42)
    
    processor = UWBRadarProcessor()
    
    # 模拟 CIR 数据（100帧, 4接收通道, 64采样点）
    # 静态杂波 + 模拟呼吸信号
    t = np.arange(100) / 20.0  # 5秒数据
    breathing_signal = 0.1 * np.sin(2 * np.pi * 0.4 * t)  # 0.4Hz = 24次/分
    
    cir = np.random.randn(100, 4, 64) * 0.01  # 噪声
    cir += 0.5  # 静态杂波
    cir[:, :, 20] += breathing_signal.reshape(-1, 1, 1)  # 距离门20处有呼吸
    
    features = processor.process_cir(cir)
    
    print("=== UWB 雷达特征 ===")
    print(f"静态能量: {features['static_energy']:.4f}")
    print(f"运动水平: {features['motion_level']:.6f}")
    print(f"呼吸幅度: {features['breathing_amplitude']:.4f}")
    print(f"距离剖面峰值位置: 门{np.argmax(features['range_profile'][0])}")

from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import numpy as np

class CPDClassifier:
    """
    Calterah 开发套件中的 ML 模型（概念实现）
    
    区分场景：
    - empty: 空座
    - adult: 成人
    - child: 儿童
    - object: 静态物体
    """
    
    def __init__(self):
        self.model = RandomForestClassifier(
            n_estimators=100,
            max_depth=10,
            random_state=42,
        )
        self.feature_names = [
            'static_energy', 'motion_level',
            'breathing_rate', 'breathing_amplitude',
            'heart_rate', 'range_peak', 'doppler_peak',
        ]
    
    def extract_features(self, processor: UWBRadarProcessor,
                        cir_data: np.ndarray) -> np.ndarray:
        """从 CIR 数据提取特征向量"""
        features = processor.process_cir(cir_data)
        
        return np.array([[
            features['static_energy'],
            features['motion_level'],
            features.get('breathing_rate', 0),
            features['breathing_amplitude'],
            features.get('heart_rate', 0),
            np.max(features['range_profile']),
            np.max(features['doppler_profile']),
        ]])
    
    def train(self, X: np.ndarray, y: np.ndarray) -> Dict:
        """训练模型"""
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=0.2, random_state=42
        )
        self.model.fit(X_train, y_train)
        
        train_acc = self.model.score(X_train, y_train)
        test_acc = self.model.score(X_test, y_test)
        
        return {
            'train_accuracy': train_acc,
            'test_accuracy': test_acc,
            'feature_importance': dict(zip(
                self.feature_names,
                self.model.feature_importances_
            )),
        }
    
    def predict(self, X: np.ndarray) -> np.ndarray:
        """预测"""
        return self.model.predict(X)
    
    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        """概率预测"""
        return self.model.predict_proba(X)


# 测试
if __name__ == "__main__":
    np.random.seed(42)
    
    clf = CPDClassifier()
    processor = UWBRadarProcessor()
    
    # 生成模拟训练数据
    n_samples = 200
    X = []
    y = []
    labels = ['empty', 'adult', 'child', 'object']
    
    for i, label in enumerate(labels):
        for _ in range(n_samples // 4):
            cir = np.random.randn(100, 4, 64) * 0.01 + 0.5
            
            if label == 'adult':
                cir[:, :, 20] += 0.3 * np.sin(2*np.pi*0.25*np.arange(100)/20).reshape(-1,1,1)
                cir[:, :, 20] += 0.1 * np.sin(2*np.pi*1.2*np.arange(100)/20).reshape(-1,1,1)
            elif label == 'child':
                cir[:, :, 15] += 0.1 * np.sin(2*np.pi*0.4*np.arange(100)/20).reshape(-1,1,1)
                cir[:, :, 15] += 0.05 * np.sin(2*np.pi*1.5*np.arange(100)/20).reshape(-1,1,1)
            elif label == 'object':
                pass  # 静态
            
            features = clf.extract_features(processor, cir)
            X.append(features[0])
            y.append(label)
    
    X = np.array(X)
    y = np.array(y)
    
    result = clf.train(X, y)
    print("=== 训练结果 ===")
    print(f"训练准确率: {result['train_accuracy']:.2%}")
    print(f"测试准确率: {result['test_accuracy']:.2%}")
    print(f"\n特征重要性:")
    for name, imp in sorted(result['feature_importance'].items(),
                            key=lambda x: -x[1]):
        print(f"  {name}: {imp:.4f}")