"""
RadarGen: 从摄像头生成雷达点云的扩散模型

简化复现框架
"""
import torch
import torch.nn as nn
import numpy as np

class RadarGen(nn.Module):
    """
    RadarGen扩散模型
    
    输入：多视角摄像头图像
    输出：合成的雷达点云
    
    核心组件：
    1. 图像VAE编码器（将摄像头图像编码到潜在空间）
    2. 雷达表示编码器（将雷达点云BEV栅格化后编码）
    3. 条件扩散模型（在潜在空间中生成雷达表示）
    4. VAE解码器（将潜在表示解码为雷达BEV图）
    """
    
    def __init__(self, n_cameras=6, image_size=224, bev_size=256):
        super().__init__()
        self.n_cameras = n_cameras
        self.image_size = image_size
        self.bev_size = bev_size
        
        # 图像VAE编码器
        self.camera_encoder = CameraVAEEncoder(
            n_views=n_cameras,
            image_size=image_size,
            latent_dim=256
        )
        
        # 雷达BEV表示
        # 点云→BEV栅格 (range, doppler, intensity)
        self.radar_channels = 3  # range, doppler, intensity
        
        # 雷达VAE
        self.radar_encoder = RadarVAEEncoder(
            bev_size=bev_size,
            n_channels=self.radar_channels,
            latent_dim=256
        )
        self.radar_decoder = RadarVAEDecoder(
            latent_dim=256,
            bev_size=bev_size,
            n_channels=self.radar_channels
        )
        
        # 条件U-Net（扩散模型核心）
        self.unet = ConditionalUNet(
            input_dim=256,  # 雷达潜在维度
            condition_dim=256,  # 摄像头潜在维度
            hidden_dim=512,
            n_steps=1000  # 扩散步数
        )
        
        # 扩散调度
        self.betas = torch.linspace(1e-4, 0.02, 1000)
        self.alphas = 1 - self.betas
        self.alpha_bars = torch.cumprod(self.alphas, dim=0)
    
    def forward(self, camera_images, radar_bev, t):
        """
        训练前向传播
        
        Args:
            camera_images: (B, n_cameras, 3, H, W) 多视角图像
            radar_bev: (B, 3, bev_size, bev_size) 雷达BEV表示
            t: (B,) 扩散时间步
        """
        # 编码摄像头图像到潜在空间
        cam_latent = self.camera_encoder(camera_images)  # (B, 256)
        
        # 编码雷达BEV到潜在空间
        radar_latent = self.radar_encoder(radar_bev)  # (B, 256)
        
        # 添加噪声
        noise = torch.randn_like(radar_latent)
        alpha_bar = self.alpha_bars[t].unsqueeze(1)
        noisy_latent = torch.sqrt(alpha_bar) * radar_latent + \
                       torch.sqrt(1 - alpha_bar) * noise
        
        # 预测噪声
        predicted_noise = self.unet(noisy_latent, cam_latent, t)
        
        return noise, predicted_noise
    
    @torch.no_grad()
    def generate(self, camera_images, n_steps=50):
        """
        从摄像头图像生成雷达点云
        
        Args:
            camera_images: (B, n_cameras, 3, H, W)
            
        Returns:
            generated_radar_bev: (B, 3, bev_size, bev_size)
        """
        # 编码摄像头条件
        cam_latent = self.camera_encoder(camera_images)
        
        # 从随机噪声开始
        x = torch.randn(cam_latent.shape[0], 256, device=cam_latent.device)
        
        # DDIM采样
        step_size = 1000 // n_steps
        for t in range(999, -1, -step_size):
            t_batch = torch.full((x.shape[0],), t, device=x.device)
            
            # 预测噪声
            predicted_noise = self.unet(x, cam_latent, t_batch)
            
            # DDIM更新
            alpha_t = self.alphas[t]
            alpha_bar_t = self.alpha_bars[t]
            
            # 去噪一步
            x0_pred = (x - torch.sqrt(1 - alpha_bar_t) * predicted_noise) / \
                       torch.sqrt(alpha_bar_t)
            x = torch.sqrt(alpha_bar_t - step_size * alpha_t) * x0_pred + \
                torch.sqrt(1 - alpha_bar_t + step_size * alpha_t) * predicted_noise
        
        # 解码为雷达BEV
        generated_radar = self.radar_decoder(x)
        
        return generated_radar


class CameraVAEEncoder(nn.Module):
    """多视角摄像头VAE编码器"""
    def __init__(self, n_views=6, image_size=224, latent_dim=256):
        super().__init__()
        self.n_views = n_views
        
        # 共享的图像编码器
        self.image_encoder = nn.Sequential(
            nn.Conv2d(3, 32, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
        )
        
        # 融合多视角
        self.fusion = nn.Sequential(
            nn.Linear(128 * n_views, 512),
            nn.ReLU(),
            nn.Linear(512, latent_dim * 2),  # mean + log_var
        )
    
    def forward(self, images):
        B = images.shape[0]
        feats = []
        for v in range(self.n_views):
            feats.append(self.image_encoder(images[:, v]))
        fused = torch.cat(feats, dim=1)
        params = self.fusion(fused)
        mean, log_var = params.chunk(2, dim=1)
        # 重参数化
        std = torch.exp(0.5 * log_var)
        eps = torch.randn_like(std)
        return mean + eps * std


class ConditionalUNet(nn.Module):
    """条件U-Net用于扩散模型"""
    def __init__(self, input_dim, condition_dim, hidden_dim=512, n_steps=1000):
        super().__init__()
        # 时间嵌入
        self.time_embed = nn.Embedding(n_steps, hidden_dim)
        
        # 条件投影
        self.cond_proj = nn.Linear(condition_dim, hidden_dim)
        
        # U-Net简化版
        self.net = nn.Sequential(
            nn.Linear(input_dim + hidden_dim * 2, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim),  # 预测噪声
        )
    
    def forward(self, x, condition, t):
        t_emb = self.time_embed(t)
        cond = self.cond_proj(condition)
        inp = torch.cat([x, t_emb, cond], dim=1)
        return self.net(inp)

"""
座舱版RadarGen：从DMS摄像头图像生成CPD雷达数据

应用场景：
1. 用大量DMS视觉数据扩充CPD雷达训练集
2. 在无雷达硬件时预训练CPD算法
3. 生成罕见场景雷达数据（儿童被毯子覆盖）
"""

class CabinRadarGen:
    """
    座舱版RadarGen
    
    输入：DMS/OMS摄像头图像
    输出：合成的座舱雷达BEV图
    
    与汽车外部RadarGen的差异：
    1. 场景从户外道路→室内座舱
    2. 距离从100m+→5m以内
    3. 目标从车辆/行人→乘员/儿童/宠物
    4. 分辨率需求更高（CPD需要10cm级）
    """
    
    def __init__(self):
        self.config = {
            'input_cameras': ['DMS_front', 'OMS_rear_left', 'OMS_rear_right'],
            'output_bev_size': 64,   # 座舱BEV 64×64（5m范围）
            'output_channels': 3,     # range, micro_doppler, intensity
            'radar_type': 'UWB_60GHz',
            'resolution': 0.05,        # 5cm分辨率
        }
        
        self.model = RadarGen(
            n_cameras=3,
            image_size=128,    # 座舱摄像头分辨率可低
            bev_size=64        # 座舱BEV较小
        )
    
    def generate_cpd_data(self, camera_images):
        """
        从摄像头生成合成CPD雷达数据
        
        用于：
        1. 扩充CPD训练数据
        2. 罕见场景生成（毯子覆盖儿童）
        3. 无雷达硬件的算法预训练
        """
        synthetic_radar = self.model.generate(camera_images)
        
        # 后处理：BEV→点云
        point_cloud = self._bev_to_pointcloud(synthetic_radar)
        
        return {
            'bev': synthetic_radar,
            'point_cloud': point_cloud,
            'quality_score': self._assess_quality(synthetic_radar),
        }
    
    def _bev_to_pointcloud(self, bev):
        """BEV栅格→雷达点云"""
        points = []
        for i in range(bev.shape[2]):
            for j in range(bev.shape[3]):
                intensity = bev[0, 2, i, j]  # intensity通道
                if intensity > 0.1:  # 有效回波
                    points.append({
                        'x': (j - 32) * 0.05,  # 5cm栅格
                        'y': (i - 32) * 0.05,
                        'intensity': intensity,
                        'doppler': bev[0, 1, i, j],
                    })
        return points

"""
合成雷达数据验证框架

三层验证：
1. 统计分布对比（合成 vs 真实）
2. 下游任务性能（用合成数据训练的CPD模型在真实数据上的表现）
3. 物理一致性（多普勒/距离/角度是否物理合理）
"""

class SyntheticRadarValidator:
    def validate(self, synthetic_data, real_data):
        results = {}
        
        # Layer 1: 统计分布
        results['distribution_distance'] = self._compute_distribution_distance(
            synthetic_data, real_data
        )
        
        # Layer 2: 下游任务
        model_synthetic = train_cpd_model(synthetic_data)
        real_accuracy = evaluate(model_synthetic, real_data)
        results['downstream_accuracy'] = real_accuracy
        
        # Layer 3: 物理一致性
        results['physics_consistency'] = self._check_physics(synthetic_data)
        
        # 通过标准
        passed = (
            results['distribution_distance'] < 0.15 and
            results['downstream_accuracy'] > 0.85 and
            results['physics_consistency'] > 0.90
        )
        
        return {'passed': passed, 'metrics': results}