import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Voxel-grid pooling for 3-D molecular point clouds.  Discretizes
    continuous 3-D coordinates into a regular grid and pools per-atom
    features into each voxel.  Used in 3-D-CNN approaches for molecular
    property prediction and ligand docking (e.g., PointVoxel-CNN,
    3D-Infomax).
    """

    def __init__(self, feat_dim, grid_size=32, voxel_size=1.0):
        super().__init__()
        self.feat_dim = feat_dim
        self.grid_size = grid_size
        self.voxel_size = voxel_size
        self.linear = nn.Linear(feat_dim, feat_dim)

    def forward(
        self, coords: torch.Tensor, features: torch.Tensor
    ) -> torch.Tensor:
        """
        Args:
            coords:   (B, N, 3) – atom coordinates
            features: (B, N, feat_dim) – per-atom features
        Returns:
            voxel_grid: (B, feat_dim, G, G, G) – pooled voxel features
        """
        B, N, _ = coords.shape
        G = self.grid_size

        features = self.linear(features)

        # Center coordinates and discretize
        center = coords.mean(dim=1, keepdim=True)
        shifted = (coords - center) / self.voxel_size + G / 2.0
        indices = shifted.long().clamp(0, G - 1)  # (B, N, 3)

        voxel_grid = coords.new_zeros(B, self.feat_dim, G, G, G)
        ix = indices[:, :, 0]
        iy = indices[:, :, 1]
        iz = indices[:, :, 2]

        batch_idx = torch.arange(B, device=coords.device).unsqueeze(1).expand(-1, N)
        for d in range(self.feat_dim):
            feat_d = features[:, :, d]
            voxel_grid[:, d].index_put_(
                (batch_idx.reshape(-1), ix.reshape(-1), iy.reshape(-1), iz.reshape(-1)),
                feat_d.reshape(-1),
                accumulate=True,
            )

        return voxel_grid


feat_dim = 32
grid_size = 32
voxel_size = 1.0
num_atoms = 256
batch_size = 4


def get_inputs():
    p = popcorn_pri
    bs = p.jitter_int(batch_size)
    na = p.jitter_int(num_atoms)
    coords = torch.randn(bs, na, 3) * 5.0
    features = torch.randn(bs, na, feat_dim)
    return [coords, features]


def get_init_inputs():
    return [feat_dim, grid_size, voxel_size]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

# Define the CUDA kernel with better parallelism
cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

// Optimized kernel: each thread handles one (atom, feature_dim) pair
// This gives B * N * feat_dim total threads for maximum parallelism
template <typename scalar_t>
__global__ void voxel_pool_kernel(
    const scalar_t* __restrict__ features,        // (B, N, feat_dim)
    const scalar_t* __restrict__ coords,          // (B, N, 3)
    scalar_t* __restrict__ voxel_grid,            // (B, feat_dim, G, G, G)
    const scalar_t* __restrict__ centers,         // (B, 3)
    int B, int N, int feat_dim, int G, scalar_t voxel_size, scalar_t half_G
) {
    const int total_work = B * N * feat_dim;
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (idx >= total_work) return;
    
    // Decode indices
    const int d = idx % feat_dim;
    const int remainder = idx / feat_dim;
    const int batch_idx = remainder / N;
    const int atom_idx = remainder % N;
    
    // Load center for this batch
    const scalar_t cx = centers[batch_idx * 3 + 0];
    const scalar_t cy = centers[batch_idx * 3 + 1];
    const scalar_t cz = centers[batch_idx * 3 + 2];
    
    // Load coordinates
    const int coord_base = batch_idx * N * 3 + atom_idx * 3;
    const scalar_t x = coords[coord_base + 0];
    const scalar_t y = coords[coord_base + 1];
    const scalar_t z = coords[coord_base + 2];
    
    // Compute grid indices
    const int ix = max(0, min(G - 1, static_cast<int>((x - cx) / voxel_size + half_G)));
    const int iy = max(0, min(G - 1, static_cast<int>((y - cy) / voxel_size + half_G)));
    const int iz = max(0, min(G - 1, static_cast<int>((z - cz) / voxel_size + half_G)));
    
    // Load feature value
    const scalar_t feat_val = features[batch_idx * N * feat_dim + atom_idx * feat_dim + d];
    
    // Compute voxel index: (B, feat_dim, G, G, G) - index is [b, d, ix, iy, iz]
    const int voxel_idx = batch_idx * feat_dim * G * G * G + d * G * G * G + ix * G * G + iy * G + iz;
    
    // Atomic add
    atomicAdd(&voxel_grid[voxel_idx], feat_val);
}

// Compute centers kernel
template <typename scalar_t>
__global__ void compute_centers_kernel(
    const scalar_t* __restrict__ coords,  // (B, N, 3)
    scalar_t* __restrict__ centers,       // (B, 3)
    int B, int N
) {
    const int batch_idx = blockIdx.x * blockDim.y + threadIdx.y;
    const int dim = threadIdx.x;  // 0, 1, or 2
    
    if (batch_idx >= B) return;
    
    scalar_t sum = 0;
    for (int i = 0; i < N; i++) {
        sum += coords[batch_idx * N * 3 + i * 3 + dim];
    }
    centers[batch_idx * 3 + dim] = sum / static_cast<scalar_t>(N);
}

torch::Tensor voxel_pool_cuda(
    torch::Tensor features,       // (B, N, feat_dim) - already transformed
    torch::Tensor coords,         // (B, N, 3)
    torch::Tensor voxel_grid,     // (B, feat_dim, G, G, G) - pre-initialized zeros
    int G, float voxel_size
) {
    const int B = coords.size(0);
    const int N = coords.size(1);
    const int feat_dim = features.size(2);
    const float half_G = static_cast<float>(G) / 2.0f;
    
    // Compute centers on GPU
    auto centers = torch::empty({B, 3}, coords.options());
    dim3 center_threads(3, 128);  // 3 dims per batch, 128 batches per block
    dim3 center_blocks((B + 127) / 128);
    
    AT_DISPATCH_FLOATING_TYPES(features.scalar_type(), "compute_centers_kernel", ([&] {
        compute_centers_kernel<scalar_t><<<center_blocks, center_threads>>>(
            coords.data_ptr<scalar_t>(),
            centers.data_ptr<scalar_t>(),
            B, N
        );
    }));
    
    const int total_work = B * N * feat_dim;
    const int threads = 256;
    const int blocks = (total_work + threads - 1) / threads;
    
    AT_DISPATCH_FLOATING_TYPES(features.scalar_type(), "voxel_pool_kernel", ([&] {
        voxel_pool_kernel<scalar_t><<<blocks, threads>>>(
            features.data_ptr<scalar_t>(),
            coords.data_ptr<scalar_t>(),
            voxel_grid.data_ptr<scalar_t>(),
            centers.data_ptr<scalar_t>(),
            B, N, feat_dim, G, static_cast<scalar_t>(voxel_size), static_cast<scalar_t>(half_G)
        );
    }));
    
    return voxel_grid;
}
'''

cpp_source = r'''
torch::Tensor voxel_pool_cuda(torch::Tensor features, torch::Tensor coords, torch::Tensor voxel_grid, int G, float voxel_size);
'''

voxel_extension = load_inline(
    name='voxel_pool_cuda_v4',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['voxel_pool_cuda'],
    extra_cuda_cflags=['-O3', '-use_fast_math'],
    verbose=True
)


class ModelNew(nn.Module):
    def __init__(self, feat_dim, grid_size=32, voxel_size=1.0):
        super().__init__()
        self.feat_dim = feat_dim
        self.grid_size = grid_size
        self.voxel_size = voxel_size
        self.linear = nn.Linear(feat_dim, feat_dim)

    def forward(self, coords: torch.Tensor, features: torch.Tensor) -> torch.Tensor:
        B, N, _ = coords.shape
        G = self.grid_size

        features = self.linear(features)

        # Initialize voxel grid
        voxel_grid = coords.new_zeros(B, self.feat_dim, G, G, G)

        # Call CUDA kernel
        voxel_grid = voxel_extension.voxel_pool_cuda(
            features.contiguous(),
            coords.contiguous(),
            voxel_grid,
            G,
            self.voxel_size
        )

        return voxel_grid