import numpy as np
from dataclasses import dataclass

@dataclass
class Radar4DData:
    """4D 雷达数据表示"""
    
    # 原始 ADC 数据
    raw_adc: np.ndarray  # (n_chirps, n_rx, n_samples)
    
    # 距离-多普勒图（Range-Doppler Map）
    range_doppler: np.ndarray  # (n_range, n_doppler, n_rx)
    
    # 距离-方位角图（Range-Angle Map）
    range_angle: np.ndarray  # (n_range, n_azimuth, n_elevation)
    
    # 点云（Point Cloud）
    points: np.ndarray  # (n_points, 6): [x, y, z, doppler, intensity, elevation]
    
    # 占用网格（Occupancy Grid）
    occupancy: np.ndarray  # (grid_x, grid_y, grid_z)
    
    # 微多普勒谱（Micro-Doppler Spectrogram）
    micro_doppler: np.ndarray  # (n_frames, n_doppler_bins)


def range_doppler_fft(raw_adc: np.ndarray) -> np.ndarray:
    """
    从原始 ADC 数据生成 Range-Doppler Map
    
    Args:
        raw_adc: shape=(n_chirps, n_rx, n_samples)
    
    Returns:
        rd_map: Range-Doppler Map, shape=(n_range, n_doppler, n_rx)
    """
    # Step 1: Range FFT（沿快时间维度）
    range_fft = np.fft.fft(raw_adc, axis=-1)
    
    # Step 2: Doppler FFT（沿慢时间/chirp 维度）
    rd_map = np.fft.fft(range_fft, axis=0)
    rd_map = np.fft.fftshift(rd_map, axes=0)
    
    # 取幅度
    rd_map = np.abs(rd_map)
    
    return rd_map


def generate_point_cloud(rd_map: np.ndarray, n_tx: int = 6, n_rx: int = 6) -> np.ndarray:
    """
    从 Range-Doppler Map 生成 4D 点云
    
    通过 MIMO 虚拟阵列进行角度估计（DBF 或 MUSIC）
    
    Args:
        rd_map: shape=(n_range, n_doppler, n_rx)
        n_tx: 发射天线数
        n_rx: 接收天线数
    
    Returns:
        points: (n_points, 6) [x, y, z, doppler, intensity, elevation]
    """
    n_virtual = n_tx * n_rx
    n_range, n_doppler, _ = rd_map.shape
    
    # CFAR 检测（简化版）
    threshold = np.mean(rd_map) + 3 * np.std(rd_map)
    detections = np.where(rd_map[:, :, 0] > threshold)
    
    points = []
    for r_idx, d_idx in zip(*detections):
        # 角度估计（简化：均匀线阵 DBF）
        azimuth = 0  # 实际应用 MUSIC/ESPRIT
        elevation = 0
        
        # 转换为笛卡尔坐标
        range_val = r_idx * 0.04  # 4cm 分辨率
        x = range_val * np.cos(elevation) * np.cos(azimuth)
        y = range_val * np.cos(elevation) * np.sin(azimuth)
        z = range_val * np.sin(elevation)
        
        doppler = (d_idx - n_doppler//2) * 0.1  # m/s
        intensity = rd_map[r_idx, d_idx, 0]
        
        points.append([x, y, z, doppler, intensity, elevation])
    
    return np.array(points)


# 测试数据生成
if __name__ == "__main__":
    np.random.seed(42)
    
    # 模拟 4D 雷达原始数据
    n_chirps = 128
    n_rx = 6
    n_samples = 256
    
    raw_adc = (np.random.randn(n_chirps, n_rx, n_samples) * 0.1 + 
               np.sin(2 * np.pi * 0.1 * np.arange(n_samples))[None, None, :] * 5)
    
    # 生成 Range-Doppler Map
    rd_map = range_doppler_fft(raw_adc)
    print(f"Range-Doppler Map shape: {rd_map.shape}")
    print(f"  Range bins: {rd_map.shape[0]}")
    print(f"  Doppler bins: {rd_map.shape[1]}")
    print(f"  RX channels: {rd_map.shape[2]}")
    
    # 生成点云
    points = generate_point_cloud(rd_map, n_tx=6, n_rx=6)
    print(f"\n检测到点云: {len(points)} 个点")
    if len(points) > 0:
        print(f"  距离范围: {np.sqrt(points[:,0]**2 + points[:,1]**2).min():.2f} - {np.sqrt(points[:,0]**2 + points[:,1]**2).max():.2f} m")
        print(f"  速度范围: {points[:,3].min():.2f} - {points[:,3].max():.2f} m/s")

import torch
import torch.nn as nn

class RadarCameraDistillation(nn.Module):
    """
    摄像头→雷达 跨模态知识蒸馏
    
    论文 Section 5.3 重点讨论：
    用高精度摄像头模型监督雷达模型训练
    """
    
    def __init__(self, radar_encoder: nn.Module, camera_encoder: nn.Module):
        super().__init__()
        self.radar_encoder = radar_encoder
        self.camera_encoder = camera_encoder
        
        # 冻结摄像头编码器（作为 Teacher）
        for p in self.camera_encoder.parameters():
            p.requires_grad = False
    
    def forward(self, radar_input, camera_input):
        """
        Args:
            radar_input: (B, C_r, H, W) 雷达 BEV 特征
            camera_input: (B, C_c, H, W) 摄像头 BEV 特征
        
        Returns:
            loss: 蒸馏损失
        """
        # Teacher: 摄像头特征
        with torch.no_grad():
            cam_feat = self.camera_encoder(camera_input)
        
        # Student: 雷达特征
        radar_feat = self.radar_encoder(radar_input)
        
        # 特征对齐（通道投影）
        if radar_feat.shape[1] != cam_feat.shape[1]:
            projector = nn.Conv2d(
                radar_feat.shape[1], cam_feat.shape[1], 1
            ).to(radar_feat.device)
            radar_feat = projector(radar_feat)
        
        # KL 散度蒸馏损失
        T = 4.0  # 温度参数
        loss_kd = nn.KLDivLoss()(
            torch.log_softmax(radar_feat / T, dim=1),
            torch.softmax(cam_feat / T, dim=1)
        ) * (T * T)
        
        return loss_kd


# 测试蒸馏框架
if __name__ == "__main__":
    # 简化编码器
    radar_enc = nn.Sequential(
        nn.Conv2d(1, 32, 3, padding=1),
        nn.ReLU(),
        nn.AdaptiveAvgPool2d((64, 64))
    )
    cam_enc = nn.Sequential(
        nn.Conv2d(3, 64, 3, padding=1),
        nn.ReLU(),
        nn.AdaptiveAvgPool2d((64, 64))
    )
    
    model = RadarCameraDistillation(radar_enc, cam_enc)
    
    radar_in = torch.randn(2, 1, 64, 64)
    cam_in = torch.randn(2, 3, 64, 64)
    
    loss = model(radar_in, cam_in)
    print(f"蒸馏损失: {loss.item():.4f}")

# 座舱雷达数据采集框架设计
class CabinRadarDataCollector:
    """座舱 4D/60GHz 雷达数据采集器"""
    
    SCENARIOS = [
        # CPD 场景
        ('cpd_empty_rear_left', '后排左座无人', 100),
        ('cpd_child_sleeping', '儿童后排睡眠', 50),
        ('cpd_child_covered', '儿童被毛毯覆盖', 50),
        ('cpd_infant_seat', '婴儿 rear-facing 座椅', 50),
        ('cpd_object_bag', '座椅放置书包', 50),
        # OOP 场景
        ('oop_normal_sitting', '正常坐姿', 100),
        ('oop_leaning_forward', '前倾', 50),
        ('oop_slouching', '瘫坐', 50),
        ('oop_feet_on_dash', '脚搭仪表台', 50),
        ('oop_unbuckled', '未系安全带', 50),
    ]
    
    def collect(self):
        for name, desc, n_samples in self.SCENARIOS:
            print(f"采集: {desc} ({name}), 目标 {n_samples} 帧")
            # 采集逻辑...