"""
RadarMind 座舱手势识别仿真

mmWave雷达手势识别核心：
1. 微多普勒谱图: 从雷达回声中提取时频表示
2. 深度学习分类: CNN/LSTM对谱图分类
3. 多手势映射: 不同运动模式对应不同手势

硬件: TI IWR1443 (77-81 GHz, 4发4收)
配置: 1帧/50ms, 256 chirps/帧
"""

import numpy as np
from scipy.signal import stft
from dataclasses import dataclass
from typing import List, Tuple
from enum import Enum

class GestureType(Enum):
    """座舱手势类型"""
    WAVE_LEFT = "wave_left"       # 左挥手
    WAVE_RIGHT = "wave_right"     # 右挥手
    ROTATE_CW = "rotate_cw"      # 顺时针旋转
    ROTATE_CCW = "rotate_ccw"    # 逆时针旋转
    PUSH = "push"                 # 前推
    PULL = "pull"                 # 后拉
    SNAP = "snap"                 # 打响指
    IDLE = "idle"                 # 无动作

@dataclass
class RadarConfig:
    """mmWave雷达配置"""
    center_freq: float = 79e9      # 79 GHz
    bandwidth: float = 4e9        # 4 GHz
    num_chirps: int = 256          # 每帧chirp数
    chirp_rate: float = 80e12      # 80 THz/s
    sample_rate: float = 8e6       # 8 Msps
    samples_per_chirp: int = 64    # 每chirp采样数
    frame_rate: float = 20         # Hz (50ms/帧)
    
    @property
    def wavelength(self):
        return 3e8 / self.center_freq  # ~3.8 mm
    
    @property
    def range_resolution(self):
        return 3e8 / (2 * self.bandwidth)  # ~3.75 cm


class MicroDopplerGenerator:
    """
    生成不同手势的微多普勒信号
    
    核心原理：
    - 手部运动 → 雷达回波相位变化
    - 相位变化率 → 多普勒频率
    - 时频分析 → 微多普勒谱图
    """
    
    def __init__(self, config: RadarConfig):
        self.config = config
        self.wavelength = config.wavelength
    
    def generate_gesture_signal(
        self, gesture: GestureType, 
        duration_s: float = 2.0
    ) -> np.ndarray:
        """
        生成手势的微多普勒信号
        
        Returns:
            spectrogram: 微多普勒谱图 [freq_bins, time_bins]
        """
        t = np.arange(0, duration_s, 1/self.config.frame_rate)
        
        # 不同手势的运动模式
        if gesture == GestureType.WAVE_LEFT:
            # 水平挥手: 正弦运动
            velocity = 0.5 * np.sin(2 * np.pi * 2 * t)  # 2Hz
        elif gesture == GestureType.WAVE_RIGHT:
            # 反向挥手
            velocity = -0.5 * np.sin(2 * np.pi * 2 * t)
        elif gesture == GestureType.ROTATE_CW:
            # 顺时针旋转: 旋转速度变化
            velocity = 0.3 * np.sin(2 * np.pi * 1.5 * t + np.pi/2)
        elif gesture == GestureType.ROTATE_CCW:
            velocity = -0.3 * np.sin(2 * np.pi * 1.5 * t + np.pi/2)
        elif gesture == GestureType.PUSH:
            # 前推: 减速运动
            velocity = 0.8 * np.exp(-t * 1.5)
        elif gesture == GestureType.PULL:
            # 后拉: 加速运动
            velocity = -0.3 * (1 - np.exp(-t * 1.5))
        elif gesture == GestureType.SNAP:
            # 打响指: 快速脉冲
            velocity = np.zeros_like(t)
            pulse_idx = int(len(t) * 0.5)
            velocity[pulse_idx:pulse_idx+2] = 2.0  # 快速脉冲
        else:
            velocity = np.zeros_like(t)
        
        # 位移 → 相位 → 信号
        displacement = np.cumsum(velocity) / self.config.frame_rate
        phase = 4 * np.pi * displacement / self.wavelength
        
        # 基带信号
        signal = np.exp(1j * phase)
        
        # 添加噪声
        noise = np.random.normal(0, 0.1, len(t)) + \
                1j * np.random.normal(0, 0.1, len(t))
        signal += noise
        
        # STFT → 微多普勒谱图
        freqs, times, spectrogram = stft(
            signal, 
            fs=self.config.frame_rate,
            nperseg=32, noverlap=28
        )
        
        return np.abs(spectrogram)
    
    def generate_dataset(
        self, num_samples: int = 100
    ) -> Tuple[np.ndarray, np.ndarray]:
        """生成手势数据集"""
        gestures = list(GestureType)
        X, y = [], []
        
        for _ in range(num_samples):
            gesture = np.random.choice(gestures)
            spec = self.generate_gesture_signal(gesture)
            
            # 标准化谱图大小
            target_shape = (17, 17)  # 标准化
            spec_resized = self._resize_spectrogram(spec, target_shape)
            
            X.append(spec_resized)
            y.append(gestures.index(gesture))
        
        return np.array(X), np.array(y)
    
    def _resize_spectrogram(
        self, spec: np.ndarray, target: Tuple[int, int]
    ) -> np.ndarray:
        """调整谱图大小"""
        from scipy.interpolate import RegularGridInterpolator
        
        h, w = spec.shape
        x = np.linspace(0, 1, w)
        y = np.linspace(0, 1, h)
        
        new_x = np.linspace(0, 1, target[1])
        new_y = np.linspace(0, 1, target[0])
        
        interp = RegularGridInterpolator((y, x), spec)
        Y, X = np.meshgrid(new_y, new_x, indexing='ij')
        return interp(np.stack([Y.ravel(), X.ravel()], axis=1)).reshape(target)


# 简化分类器
class GestureClassifier:
    """基于谱图的手势分类器"""
    
    def __init__(self, num_classes: int = 8):
        from sklearn.ensemble import RandomForestClassifier
        self.model = RandomForestClassifier(n_estimators=100, random_state=42)
        self.num_classes = num_classes
        self.trained = False
    
    def fit(self, X: np.ndarray, y: np.ndarray):
        X_flat = X.reshape(X.shape[0], -1)
        self.model.fit(X_flat, y)
        self.trained = True
    
    def predict(self, X: np.ndarray) -> np.ndarray:
        X_flat = X.reshape(X.shape[0], -1)
        return self.model.predict(X_flat)


# ==================== 测试 ====================
if __name__ == "__main__":
    config = RadarConfig()
    gen = MicroDopplerGenerator(config)
    
    print("=== RadarMind 座舱手势识别仿真 ===")
    print(f"雷达: TI IWR1443")
    print(f"频段: {config.center_freq/1e9:.0f} GHz")
    print(f"波长: {config.wavelength*1000:.1f} mm")
    print(f"距离分辨率: {config.range_resolution*100:.1f} cm")
    print(f"帧率: {config.frame_rate} Hz")
    
    # 生成数据集
    X, y = gen.generate_dataset(num_samples=200)
    print(f"\n数据集: {X.shape[0]} 样本, {X.shape[1]}x{X.shape[2]} 谱图")
    
    # 分类
    clf = GestureClassifier()
    
    # 训练/测试分割
    split = int(0.8 * len(X))
    X_train, X_test = X[:split], X[split:]
    y_train, y_test = y[:split], y[split:]
    
    clf.fit(X_train, y_train)
    accuracy = (clf.predict(X_test) == y_test).mean()
    
    print(f"\n分类准确率: {accuracy*100:.1f}%")
    
    # 各手势示例
    print(f"\n--- 手势微多普勒特征 ---")
    for gesture in GestureType:
        spec = gen.generate_gesture_signal(gesture)
        print(f"  {gesture.value:15s}: 谱图峰值={np.max(spec):.3f}, "
              f"频率展宽={np.std(np.where(spec > 0.1*np.max(spec), 1, 0)):.3f}")