import numpy as np
import torch
import torch.nn as nn
from scipy import signal as sig

class RadarGestureRecognition:
    """mmWave 雷达手势识别系统"""
    
    def __init__(self, n_classes=8):
        self.n_classes = n_classes
        
        # IWR1443参数
        self.n_rx = 4          # 接收天线数
        self.n_chirp = 64      # 每帧chirp数
        self.n_adc = 256       # 每chirp ADC采样数
        self.frame_rate = 30   # fps
        
        # 信号处理参数
        self.range_fft_size = 256
        self.doppler_fft_size = 64
        self.angle_fft_size = 64
        
        # 深度学习模型
        self.model = RadarGestureNet(n_classes=n_classes)
        
    def process_raw_data(self, adc_data):
        """
        处理IWR1443原始ADC数据
        
        Args:
            adc_data: (n_chirp, n_rx, n_adc) raw ADC
        
        Returns:
            range_doppler_map: (n_rx, range_bins, doppler_bins)
            range_angle_map: (doppler_bins, range_bins, angle_bins)
        """
        # 1. Range FFT (沿ADC维度)
        range_fft = np.fft.fft(adc_data, n=self.range_fft_size, axis=-1)
        range_fft = np.abs(range_fft)  # (n_chirp, n_rx, range_bins)
        
        # 2. Doppler FFT (沿chirp维度)
        doppler_fft = np.fft.fftshift(
            np.fft.fft(range_fft, n=self.doppler_fft_size, axis=0),
            axes=0
        )  # (doppler_bins, n_rx, range_bins)
        
        # 3. Angle FFT (沿rx维度)
        angle_fft = np.fft.fftshift(
            np.fft.fft(doppler_fft, n=self.angle_fft_size, axis=1),
            axes=1
        )  # (doppler_bins, angle_bins, range_bins)
        
        # 取log幅度
        range_doppler = np.log1p(np.abs(doppler_fft))  # (doppler, n_rx, range)
        range_angle = np.log1p(np.abs(angle_fft))     # (doppler, angle, range)
        
        return range_doppler, range_angle
    
    def extract_micro_doppler(self, range_doppler, n_frames=30):
        """
        提取micro-Doppler特征（时频图）
        
        手势的微多普勒特征包含：
        - 手部运动速度
        - 运动方向变化
        - 运动持续时间
        """
        # 取中心距离bin的时间序列
        center_bin = range_doppler.shape[-1] // 2
        
        # 对每帧取最大速度
        max_velocity = np.max(np.abs(range_doppler[:, :, center_bin]), axis=0)
        
        # STFT获取时频图
        f, t, sxx = sig.stft(
            max_velocity, 
            fs=self.frame_rate,
            nperseg=8,
            noverlap=4
        )
        
        return sxx  # (freq_bins, time_bins)


class RadarGestureNet(nn.Module):
    """mmWave雷达手势分类网络"""
    
    def __init__(self, n_classes=8, input_channels=4):
        super().__init__()
        
        # 输入: (B, 4, 64, 64) - 4个RX的Range-Doppler图
        self.features = nn.Sequential(
            # Block 1
            nn.Conv2d(input_channels, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),  # 32x32
            
            # Block 2
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),  # 16x16
            
            # Block 3
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2),  # 8x8
            
            # Block 4
            nn.Conv2d(128, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        
        # 时序建模
        self.temporal = nn.LSTM(
            input_size=256,
            hidden_size=128,
            num_layers=2,
            batch_first=True
        )
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64, n_classes)
        )
        
    def forward(self, x):
        """
        Args:
            x: (B, T, C, H, W) T帧的Range-Doppler图
        
        Returns:
            logits: (B, n_classes)
        """
        B, T, C, H, W = x.shape
        
        # 每帧空间特征
        x = x.view(B * T, C, H, W)
        spatial = self.features(x)  # (B*T, 256)
        spatial = spatial.view(B, T, -1)  # (B, T, 256)
        
        # 时序建模
        temporal_out, _ = self.temporal(spatial)  # (B, T, 128)
        last = temporal_out[:, -1, :]  # (B, 128) 取最后一帧
        
        # 分类
        return self.classifier(last)


# === 座舱手势集定义 ===
CABIN_GESTURES = {
    0: 'swipe_left',       # 左滑（上一曲）
    1: 'swipe_right',      # 右滑（下一曲）
    2: 'swipe_up',         # 上滑（音量+）
    3: 'swipe_down',       # 下滑（音量-）
    4: 'circle_cw',        # 顺时针旋转（调温+）
    5: 'circle_ccw',       # 逆时针旋转（调温-）
    6: 'push',             # 前推（确认/接听）
    7: 'pull',             # 后拉（取消/挂断）
}

# 测试
if __name__ == "__main__":
    system = RadarGestureRecognition(n_classes=8)
    
    # 模拟ADC数据
    adc = np.random.randn(64, 4, 256) + 1j * np.random.randn(64, 4, 256)
    rd, ra = system.process_raw_data(adc)
    print(f"Range-Doppler: {rd.shape}")
    print(f"Range-Angle: {ra.shape}")
    
    # 模拟手势分类
    B, T = 2, 10  # 2个样本, 10帧
    input_data = torch.randn(B, T, 4, 64, 64)
    output = system.model(input_data)
    print(f"\n输入: {input_data.shape}")
    print(f"输出: {output.shape} (8类手势)")
    print(f"参数量: {sum(p.numel() for p in system.model.parameters())/1e6:.2f}M")
    
    print(f"\n座舱手势集:")
    for idx, name in CABIN_GESTURES.items():
        print(f"  {idx}: {name}")