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

@dataclass
class FMCWConfig:
    """60GHz FMCW 雷达配置"""
    center_freq: float = 60e9      # 60 GHz
    bandwidth: float = 4e9          # 4 GHz 扫频带宽
    sweep_duration: float = 100e-6  # 100μs 扫频时间
    sample_rate: float = 2e6        # 2 MHz 采样率
    num_chirps: int = 64            # 每帧 chirp 数
    num_tx: int = 1                 # 发射天线数
    num_rx: int = 2                 # 接收天线数

class FMCWRadarProcessor:
    """
    60GHz FMCW 雷达信号处理
    
    处理流程：
    1. ADC 采样 → IF 信号
    2. Range-FFT → 距离信息
    3. Doppler-FFT → 速度信息
    4. Angle-FFT → 角度信息
    5. CFAR → 目标检测
    6. 生命体征提取 → CPD 判定
    """
    
    def __init__(self, config: FMCWConfig):
        self.config = config
        self.range_resolution = config.bandwidth / (2 * config.center_freq)
        self.max_range = config.sample_rate * config.sweep_duration * 3e8 / (2 * config.bandwidth)
        
    def range_processing(self, adc_data: np.ndarray) -> np.ndarray:
        """
        距离处理：Range-FFT
        
        Args:
            adc_data: (num_chirps, num_samples) ADC 采样数据
            
        Returns:
            range_profile: (num_chirps, num_range_bins) 距离剖面
        """
        # 对每个 chirp 做 FFT
        range_fft = np.fft.fft(adc_data, axis=1)
        range_fft = np.fft.fftshift(range_fft, axes=1)
        
        # 取幅度
        range_profile = np.abs(range_fft)
        
        return range_profile
    
    def doppler_processing(self, range_profile: np.ndarray) -> np.ndarray:
        """
        多普勒处理：Doppler-FFT
        
        Args:
            range_profile: (num_chirps, num_range_bins) 距离剖面
            
        Returns:
            range_doppler_map: (num_doppler_bins, num_range_bins) 距离-多普勒图
        """
        # 沿 chirp 维度做 FFT
        doppler_fft = np.fft.fft(range_profile, axis=0)
        range_doppler_map = np.abs(doppler_fft)
        
        return range_doppler_map
    
    def extract_vital_signs(
        self,
        adc_data: np.ndarray,
        target_range_bin: int
    ) -> dict:
        """
        生命体征提取
        
        从特定距离 bin 提取呼吸和心跳
        
        Args:
            adc_data: 原始 ADC 数据
            target_range_bin: 目标所在距离 bin
            
        Returns:
            vital_signs: 生命体征数据
        """
        # 提取目标距离 bin 的相位变化
        range_fft = np.fft.fft(adc_data, axis=1)
        phase = np.angle(range_fft[:, target_range_bin])
        
        # 解卷绕
        phase_unwrapped = np.unwrap(phase)
        
        # 去除线性趋势（身体整体运动）
        from scipy.signal import detrend
        phase_detrended = detrend(phase_unwrapped)
        
        # 带通滤波分离呼吸和心跳
        from scipy.signal import butter, filtfilt
        
        # 呼吸频段：0.1-0.5 Hz (6-30 次/分)
        b_resp, a_resp = butter(2, [0.1, 0.5], btype='band', fs=1/self.config.sweep_duration)
        breathing = filtfilt(b_resp, a_resp, phase_detrended)
        
        # 心跳频段：0.8-2.0 Hz (48-120 次/分)
        b_heart, a_heart = butter(2, [0.8, 2.0], btype='band', fs=1/self.config.sweep_duration)
        heartbeat = filtfilt(b_heart, a_heart, phase_detrended)
        
        # FFT 频谱分析
        freqs = np.fft.fftfreq(len(phase_detrended), d=self.config.sweep_duration)
        spectrum = np.abs(np.fft.fft(phase_detrended))
        
        # 找呼吸峰值
        resp_mask = (freqs >= 0.1) & (freqs <= 0.5)
        if np.any(resp_mask):
            resp_freq = freqs[resp_mask][np.argmax(spectrum[resp_mask])]
            breathing_rate = resp_freq * 60  # 次/分
        else:
            breathing_rate = 0
        
        # 找心跳峰值
        heart_mask = (freqs >= 0.8) & (freqs <= 2.0)
        if np.any(heart_mask):
            heart_freq = freqs[heart_mask][np.argmax(spectrum[heart_mask])]
            heart_rate = heart_freq * 60  # 次/分
        else:
            heart_rate = 0
        
        return {
            'breathing_rate': breathing_rate,
            'heart_rate': heart_rate,
            'breathing_signal': breathing,
            'heartbeat_signal': heartbeat,
            'phase_variation': phase_detrended,
            'spectrum': spectrum,
            'frequencies': freqs
        }
    
    def detect_child_presence(
        self,
        adc_data: np.ndarray,
        vehicle_locked: bool = True
    ) -> dict:
        """
        CPD 儿童检测主函数
        
        检测逻辑：
        1. 车辆锁定后启动检测
        2. 扫描座舱范围 (0.3m - 2.0m)
        3. 检测微动（呼吸引起的胸部起伏）
        4. 检测生命体征频率
        5. 判定是否存在儿童
        
        Args:
            adc_data: 多帧 ADC 数据
            vehicle_locked: 车辆是否已锁定
            
        Returns:
            result: CPD 检测结果
        """
        if not vehicle_locked:
            return {'status': 'DISABLED', 'reason': '车辆未锁定'}
        
        # 距离处理
        range_profile = self.range_processing(adc_data)
        
        # 扫描各距离 bin
        detections = []
        for rbin in range(10, 50):  # 0.3m - 1.5m 范围
            vital = self.extract_vital_signs(adc_data, rbin)
            
            # 判定逻辑
            breathing = vital['breathing_rate'] > 0
            heartbeat = vital['heart_rate'] > 0
            micro_motion = np.std(vital['phase_variation']) > 0.01
            
            if breathing and micro_motion:
                # 估计体型（基于反射强度）
                reflection_strength = np.mean(range_profile[:, rbin])
                
                # 儿童判定：呼吸频率偏高 + 反射面积小
                is_child_likely = (
                    vital['breathing_rate'] > 20 and  # 儿童呼吸偏快
                    reflection_strength < 0.5        # 反射面积较小
                )
                
                detections.append({
                    'range_bin': rbin,
                    'range_m': rbin * self.range_resolution,
                    'breathing_rate': vital['breathing_rate'],
                    'heart_rate': vital['heart_rate'],
                    'is_child_likely': is_child_likely,
                    'confidence': min(
                        np.std(vital['phase_variation']) * 10, 1.0
                    )
                })
        
        if detections:
            best = max(detections, key=lambda x: x['confidence'])
            return {
                'status': 'CHILD_DETECTED' if best['is_child_likely'] else 'OCCUPANT_DETECTED',
                'range_m': best['range_m'],
                'breathing_rate': best['breathing_rate'],
                'heart_rate': best['heart_rate'],
                'confidence': best['confidence'],
                'all_detections': detections
            }
        
        return {'status': 'EMPTY', 'confidence': 0.95}


if __name__ == "__main__":
    config = FMCWConfig()
    processor = FMCWRadarProcessor(config)
    
    print(f"=== 60GHz FMCW 雷达配置 ===")
    print(f"中心频率: {config.center_freq/1e9:.0f} GHz")
    print(f"带宽: {config.bandwidth/1e9:.1f} GHz")
    print(f"距离分辨率: {processor.range_resolution*100:.1f} cm")
    print(f"最大探测距离: {processor.max_range:.1f} m")
    print(f"Chirp 数: {config.num_chirps}")
    print(f"扫频时间: {config.sweep_duration*1e6:.0f} μs")
    
    # 模拟数据测试
    np.random.seed(42)
    adc_data = np.random.randn(config.num_chirps, 256) * 0.1
    
    # 添加模拟呼吸信号 (0.3Hz = 18次/分)
    t = np.arange(256) * config.sweep_duration
    breathing_signal = 0.5 * np.sin(2 * np.pi * 0.3 * t)
    adc_data[32, :] += breathing_signal  # 在某个距离 bin 注入
    
    result = processor.detect_child_presence(adc_data)
    print(f"\n=== CPD 检测结果 ===")
    print(f"状态: {result['status']}")
    if 'breathing_rate' in result:
        print(f"呼吸频率: {result['breathing_rate']:.1f} 次/分")
        print(f"心跳频率: {result['heart_rate']:.1f} 次/分")
        print(f"置信度: {result['confidence']:.2f}")
        print(f"距离: {result['range_m']:.2f} m")
    
    print("\n=== 三大方案关键指标 ===")
    print(f"{'指标':<20} {'Infineon':<15} {'TI':<15} {'NOVELIC':<15}")
    print(f"{'芯片型号':<20} {'BGT60ATR24':<15} {'AWRL6432':<15} {'ACAM':<15}")
    print(f"{'功耗(睡眠)':<20} {'<1mW':<15} {'<1mW':<15} {'<5mW':<15}")
    print(f"{'检测距离':<20} {'1.5m':<15} {'2.0m':<15} {'3.0m':<15}")
    print(f"{'距离分辨率':<20} {'~5cm':<15} {'~4cm':<15} {'~3cm':<15}")
    print(f"{'ASIL等级':<20} {'ASIL-B':<15} {'ASIL-B':<15} {'ASIL-A':<15}")
    print(f"{'心跳检测':<20} {'❌':<15} {'✅':<15} {'✅':<15}")
    print(f"{'预估价格':<20} {'$3-5':<15} {'$4-6':<15} {'$5-8':<15}")

class CPDModule:
    """
    IMS CPD 模块 - 集成到座舱监控系统
    
    与 DMS/OMS 共享：
    - 车辆状态（锁止/行驶/停车）
    - 座椅状态（占用/空置）
    - 告警系统
    
    独立运行：
    - 60GHz 雷达数据处理
    - 生命体征检测
    - CPD 判定逻辑
    """
    
    def __init__(self, config: dict = None):
        config = config or {}
        self.radar_config = FMCWConfig()
        self.processor = FMCWRadarProcessor(self.radar_config)
        
        # 状态机
        self.state = 'IDLE'  # IDLE/SCANNING/DETECTED/ALERT
        self.vehicle_locked = False
        self.detection_count = 0
        self.false_alarm_count = 0
        
        # 配置
        self.scan_interval = config.get('scan_interval', 5)  # 秒
        self.alert_threshold = config.get('alert_threshold', 0.7)
        self.confirmation_count = config.get('confirmation_count', 3)
        
    def on_vehicle_lock(self):
        """车辆锁定事件"""
        self.state = 'SCANNING'
        self.detection_count = 0
        
    def on_vehicle_unlock(self):
        """车辆解锁事件"""
        self.state = 'IDLE'
        self.detection_count = 0
        
    def process_frame(self, adc_data: np.ndarray) -> dict:
        """处理一帧雷达数据"""
        if self.state == 'IDLE':
            return {'state': 'IDLE'}
        
        result = self.processor.detect_child_presence(
            adc_data, vehicle_locked=True
        )
        
        if result['status'] == 'CHILD_DETECTED':
            self.detection_count += 1
            if self.detection_count >= self.confirmation_count:
                self.state = 'ALERT'
                return {
                    'state': 'ALERT',
                    'action': 'TRIGGER_ALARM',
                    'confidence': result.get('confidence', 0),
                    'breathing_rate': result.get('breathing_rate', 0),
                    'range_m': result.get('range_m', 0)
                }
        elif result['status'] == 'OCCUPANT_DETECTED':
            self.detection_count += 1
            return {'state': 'SCANNING', 'occupant': True}
        else:
            self.detection_count = max(0, self.detection_count - 1)
            
        return {'state': self.state, 'result': result}