import torch
import torch.nn as nn
import torch.nn.functional as F

class RadarDensifier(nn.Module):
    """
    Radar Densifier: 解决雷达点云稀疏性
    
    论文核心: 将非空雷达 "pillar" 的特征
    用自适应高斯函数重分配到邻域空 pillar
    
    高斯宽度基于 RCS 自适应:
    - 大 RCS (大物体) → 宽高斯 (远距离传播)
    - 小 RCS (小物体) → 窄高斯 (精细传播)
    """
    
    def __init__(self, in_channels=64, out_channels=64, 
                 grid_size=(200, 200), max_range=50.0):
        super().__init__()
        self.grid_size = grid_size
        self.max_range = max_range
        
        # 雷达特征编码
        self.feature_encoder = nn.Sequential(
            nn.Linear(in_channels + 5, 64),  # radar feat + [x, y, z, rcs, doppler]
            nn.ReLU(),
            nn.Linear(64, out_channels)
        )
        
        # RCS → 高斯宽度映射
        self.rcs_to_sigma = nn.Sequential(
            nn.Linear(1, 16),
            nn.ReLU(),
            nn.Linear(16, 1),
            nn.Softplus()  # 确保 > 0
        )
        
    def forward(self, radar_points, pillar_grid=None):
        """
        Args:
            radar_points: (N, 7+) [x, y, z, rcs, doppler, intensity, ...]
            pillar_grid: (G, G, C) 预初始化的 pillar 网格
            
        Returns:
            densified_grid: (G, G, C) 致密化后的网格
        """
        N = radar_points.shape[0]
        G = self.grid_size[0]
        
        # 将雷达点投影到 BEV 网格
        x = radar_points[:, 0]
        y = radar_points[:, 1]
        rcs = radar_points[:, 3]
        
        # 归一化到网格坐标
        gx = ((x + self.max_range) / (2 * self.max_range) * G).long().clamp(0, G-1)
        gy = ((y + self.max_range) / (2 * self.max_range) * G).long().clamp(0, G-1)
        
        # 初始化 pillar 网格
        grid = torch.zeros(G, G, 64)
        
        # 填充非空 pillar
        for i in range(N):
            if 0 <= gx[i] < G and 0 <= gy[i] < G:
                feat = self.feature_encoder(
                    radar_points[i:i+1]
                ).squeeze(0)
                grid[gx[i], gy[i]] += feat
        
        # 自适应高斯传播
        densified_grid = grid.clone()
        sigma_base = self.rcs_to_sigma(rcs.unsqueeze(-1)).squeeze(-1)
        
        # 对每个非空 pillar 做高斯传播
        non_empty_mask = grid.sum(dim=-1) > 0
        non_empty_idx = torch.where(non_empty_mask)
        
        for (px, py) in zip(*non_empty_idx):
            sigma = sigma_base[px * G + py] if px * G + py < N else 1.0
            sigma = max(sigma.item(), 0.5)
            
            # 高斯窗口
            radius = int(3 * sigma)
            for dx in range(-radius, radius + 1):
                for dy in range(-radius, radius + 1):
                    nx, ny = px + dx, py + dy
                    if 0 <= nx < G and 0 <= ny < G:
                        if dx == 0 and dy == 0:
                            continue
                        dist_sq = dx * dx + dy * dy
                        weight = torch.exp(
                            torch.tensor(-dist_sq / (2 * sigma ** 2))
                        )
                        densified_grid[nx, ny] += weight * grid[px, py] * 0.3
        
        return densified_grid


class RadarAmplifier(nn.Module):
    """
    Radar Amplifier: 特征通道选择性增强
    
    估计每个特征通道的重要性概率
    增强有用通道，抑制噪声通道
    """
    
    def __init__(self, channels=64, reduction=8):
        super().__init__()
        
        # 通道注意力 (类似 SE 但使用概率)
        self.channel_gate = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(channels, channels // reduction),
            nn.ReLU(),
            nn.Linear(channels // reduction, channels),
            nn.Sigmoid()  # 概率分数
        )
        
        # 通道缩放
        self.scale = nn.Parameter(torch.ones(1, channels, 1, 1))
        
    def forward(self, x):
        """
        Args:
            x: (B, C, H, W) 雷达特征图
            
        Returns:
            amplified: (B, C, H, W) 增强后的特征
        """
        # 通道重要性概率
        importance = self.channel_gate(x)  # (B, C)
        importance = importance.unsqueeze(-1).unsqueeze(-1)  # (B, C, 1, 1)
        
        # 选择性增强
        amplified = x * importance * self.scale
        
        return amplified


class REOccFusion(nn.Module):
    """
    REOcc: 摄像头-雷达融合 3D 占用预测
    
    架构:
    1. 雷达 → Densifier → Amplifier → 增强雷达特征
    2. 摄像头 → 2D backbone → BEV 转换 → 摄像头 BEV 特征
    3. 可变形交叉注意力融合
    4. 3D 占用预测头
    """
    
    def __init__(self, radar_channels=64, cam_channels=256, 
                 fusion_dim=256, num_classes=17):
        super().__init__()
        
        # 雷达处理
        self.radar_densifier = RadarDensifier(radar_channels)
        self.radar_amplifier = RadarAmplifier(radar_channels)
        self.radar_proj = nn.Conv2d(radar_channels, fusion_dim, 1)
        
        # 摄像头处理 (简化)
        self.cam_backbone = nn.Sequential(
            nn.Conv2d(3, 64, 7, stride=2, padding=3),
            nn.BatchNorm2d(64), nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm2d(128), nn.ReLU(),
            nn.Conv2d(128, cam_channels, 3, stride=2, padding=1),
            nn.BatchNorm2d(cam_channels), nn.ReLU(),
        )
        self.cam_proj = nn.Conv2d(cam_channels, fusion_dim, 1)
        
        # 可变形交叉注意力
        self.fusion = DeformableCrossAttention(
            dim=fusion_dim, num_heads=8
        )
        
        # 占用预测头
        self.occ_head = nn.Sequential(
            nn.Conv3d(fusion_dim, 128, 3, padding=1),
            nn.BatchNorm3d(128), nn.ReLU(),
            nn.Conv3d(128, num_classes, 1)
        )
        
    def forward(self, radar_points, cam_images):
        """
        Args:
            radar_points: (N, 7) 雷达点
            cam_images: (B, 3, H, W) 摄像头图像
            
        Returns:
            occ_pred: (B, num_classes, D, H, W) 3D 占用预测
        """
        # 雷达处理
        radar_grid = self.radar_densifier(radar_points)  # (G, G, C)
        radar_grid = radar_grid.permute(2, 0, 1).unsqueeze(0)  # (1, C, G, G)
        radar_feat = self.radar_amplifier(radar_grid)
        radar_feat = self.radar_proj(radar_feat)  # (B, D, G, G)
        
        # 摄像头处理
        cam_feat = self.cam_backbone(cam_images)
        cam_feat = self.cam_proj(cam_feat)  # (B, D, H', W')
        
        # 可变形交叉注意力融合
        fused = self.fusion(cam_feat, radar_feat)
        
        # 3D 占用预测 (简化: BEV → 3D)
        occ = self.occ_head(fused.unsqueeze(-1))  # 简化版
        
        return occ


class DeformableCrossAttention(nn.Module):
    """可变形交叉注意力"""
    
    def __init__(self, dim=256, num_heads=8):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        
        self.q_proj = nn.Linear(dim, dim)
        self.k_proj = nn.Linear(dim, dim)
        self.v_proj = nn.Linear(dim, dim)
        
        # 可变形偏移
        self.offset_gen = nn.Linear(dim, num_heads * 2)
        
        self.out_proj = nn.Linear(dim, dim)
        self.scale = self.head_dim ** -0.5
        
    def forward(self, query_feat, key_value_feat):
        B, C, H, W = query_feat.shape
        
        q = self.q_proj(
            query_feat.flatten(2).transpose(1, 2)
        ).transpose(1, 2).reshape(B, self.num_heads, -1, H*W)
        
        k = self.k_proj(
            key_value_feat.flatten(2).transpose(1, 2)
        ).transpose(1, 2).reshape(B, self.num_heads, -1, H*W)
        
        v = self.v_proj(
            key_value_feat.flatten(2).transpose(1, 2)
        ).transpose(1, 2).reshape(B, self.num_heads, -1, H*W)
        
        attn = torch.matmul(q.transpose(-1, -2), k) * self.scale
        attn = torch.softmax(attn, dim=-1)
        
        out = torch.matmul(attn, v.transpose(-1, -2))
        out = out.reshape(B, C, H, W)
        out = self.out_proj(
            out.flatten(2).transpose(1, 2)
        ).transpose(1, 2).reshape(B, C, H, W)
        
        return query_feat + out


if __name__ == "__main__":
    # 测试 RadarAmplifier
    amp = RadarAmplifier(64)
    x = torch.randn(2, 64, 200, 200)
    out = amp(x)
    print(f"RadarAmplifier: {x.shape} → {out.shape}")
    print(f"Params: {sum(p.numel() for p in amp.parameters()):,}")