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 CUDA source with optimized kernels
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>

// Kernel 1: Compute centers for each batch using block reduction
template <typename scalar_t>
__global__ void __launch_bounds__(256) compute_centers_kernel(
    const scalar_t* __restrict__ coords,
    scalar_t* __restrict__ centers,
    int N) {
    
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int blockSize = blockDim.x;
    
    // Use shared memory for reduction
    __shared__ scalar_t smem[3 * 256];
    
    scalar_t local_sum_x = 0.0f;
    scalar_t local_sum_y = 0.0f;
    scalar_t local_sum_z = 0.0f;
    
    // Each thread sums a subset of atoms
    for (int i = tid; i < N; i += blockSize) {
        local_sum_x += coords[b * N * 3 + i * 3 + 0];
        local_sum_y += coords[b * N * 3 + i * 3 + 1];
        local_sum_z += coords[b * N * 3 + i * 3 + 2];
    }
    
    smem[tid * 3 + 0] = local_sum_x;
    smem[tid * 3 + 1] = local_sum_y;
    smem[tid * 3 + 2] = local_sum_z;
    
    __syncthreads();
    
    // Block reduction
    for (int s = blockSize / 2; s > 0; s >>= 1) {
        if (tid < s) {
            smem[tid * 3 + 0] += smem[(tid + s) * 3 + 0];
            smem[tid * 3 + 1] += smem[(tid + s) * 3 + 1];
            smem[tid * 3 + 2] += smem[(tid + s) * 3 + 2];
        }
        __syncthreads();
    }
    
    // Thread 0 writes result
    if (tid == 0) {
        centers[b * 3 + 0] = smem[0] / N;
        centers[b * 3 + 1] = smem[1] / N;
        centers[b * 3 + 2] = smem[2] / N;
    }
}

// Kernel 2: Voxel pooling with optimized memory access
template <typename scalar_t, int FEAT_UNROLL=8>
__global__ void __launch_bounds__(256) voxel_pool_kernel(
    const scalar_t* __restrict__ coords,
    const scalar_t* __restrict__ centers,
    const scalar_t* __restrict__ features,
    scalar_t* __restrict__ voxel_grid,
    int B, int N, int G, int feat_dim,
    float voxel_size) {
    
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * N;
    
    if (idx >= total) return;
    
    int b = idx / N;
    int n = idx % N;
    
    // Get center for this batch
    scalar_t cx = centers[b * 3 + 0];
    scalar_t cy = centers[b * 3 + 1];
    scalar_t cz = centers[b * 3 + 2];
    
    // Get coordinates for this atom
    scalar_t x = coords[b * N * 3 + n * 3 + 0];
    scalar_t y = coords[b * N * 3 + n * 3 + 1];
    scalar_t z = coords[b * N * 3 + n * 3 + 2];
    
    // Compute voxel indices
    int ix = (int)floorf((x - cx) / voxel_size + G / 2.0f);
    int iy = (int)floorf((y - cy) / voxel_size + G / 2.0f);
    int iz = (int)floorf((z - cz) / voxel_size + G / 2.0f);
    
    // Clamp to valid range
    ix = max(0, min(G - 1, ix));
    iy = max(0, min(G - 1, iy));
    iz = max(0, min(G - 1, iz));
    
    // Compute linear indices
    int spatial_idx = ix * G * G + iy * G + iz;
    int base_voxel = b * feat_dim * G * G * G;
    int feat_offset = n * feat_dim;
    int batch_feat_offset = b * N * feat_dim;
    
    // Unrolled feature accumulation
    int d = 0;
    #pragma unroll
    for (; d + FEAT_UNROLL <= feat_dim; d += FEAT_UNROLL) {
        scalar_t feat[FEAT_UNROLL];
        #pragma unroll
        for (int j = 0; j < FEAT_UNROLL; j++) {
            feat[j] = features[batch_feat_offset + feat_offset + d + j];
        }
        #pragma unroll
        for (int j = 0; j < FEAT_UNROLL; j++) {
            atomicAdd(&voxel_grid[base_voxel + (d + j) * G * G * G + spatial_idx], feat[j]);
        }
    }
    
    // Handle remaining features
    for (; d < feat_dim; d++) {
        scalar_t feat = features[batch_feat_offset + feat_offset + d];
        atomicAdd(&voxel_grid[base_voxel + d * G * G * G + spatial_idx], feat);
    }
}

torch::Tensor voxel_pool_cuda(
    torch::Tensor coords,
    torch::Tensor features,
    int grid_size,
    float voxel_size) {
    
    int B = coords.size(0);
    int N = coords.size(1);
    int G = grid_size;
    int feat_dim = features.size(2);
    
    // Allocate outputs
    auto centers = torch::empty({B, 3}, coords.options());
    auto voxel_grid = torch::zeros({B, feat_dim, G, G, G}, coords.options());
    
    int threads = 256;
    
    // Kernel 1: Compute centers
    AT_DISPATCH_FLOATING_TYPES(coords.scalar_type(), "compute_centers_kernel", ([&] {
        compute_centers_kernel<scalar_t><<<B, threads>>>(
            coords.data_ptr<scalar_t>(),
            centers.data_ptr<scalar_t>(),
            N);
    }));
    
    // Kernel 2: Voxel pooling
    int blocks = (B * N + threads - 1) / threads;
    AT_DISPATCH_FLOATING_TYPES(coords.scalar_type(), "voxel_pool_kernel", ([&] {
        voxel_pool_kernel<scalar_t, 8><<<blocks, threads>>>(
            coords.data_ptr<scalar_t>(),
            centers.data_ptr<scalar_t>(),
            features.data_ptr<scalar_t>(),
            voxel_grid.data_ptr<scalar_t>(),
            B, N, G, feat_dim, voxel_size);
    }));
    
    return voxel_grid;
}
"""

cpp_source = """
torch::Tensor voxel_pool_cuda(torch::Tensor coords, torch::Tensor features, int grid_size, float voxel_size);
"""

# Load the CUDA extension
voxel_extension = load_inline(
    name='voxel_pool',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['voxel_pool_cuda'],
    verbose=False,
)


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:
        # Apply linear transformation first
        features = self.linear(features)
        
        # Call CUDA kernel for voxel pooling
        voxel_grid = voxel_extension.voxel_pool_cuda(
            coords, features, self.grid_size, self.voxel_size
        )
        
        return voxel_grid