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

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor voxel_pool_forward(torch::Tensor coords, torch::Tensor features, torch::Tensor weight, torch::Tensor bias, int64_t grid_size, double voxel_size);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void center_kernel(const float* __restrict__ coords, float* __restrict__ centers, int N) {
    int b = blockIdx.x, tid = threadIdx.x;
    float sx = 0.f, sy = 0.f, sz = 0.f;
    for (int n = tid; n < N; n += 256) {
        const float* p = coords + ((b * N + n) * 3);
        sx += p[0]; sy += p[1]; sz += p[2];
    }
    __shared__ float shx[256], shy[256], shz[256];
    shx[tid] = sx; shy[tid] = sy; shz[tid] = sz;
    __syncthreads();
    for (int s = 128; s > 0; s >>= 1) {
        if (tid < s) {
            shx[tid] += shx[tid+s];
            shy[tid] += shy[tid+s];
            shz[tid] += shz[tid+s];
        }
        __syncthreads();
    }
    if (tid == 0) {
        float inv = 1.0f / (float)N;
        centers[b*3+0] = shx[0] * inv;
        centers[b*3+1] = shy[0] * inv;
        centers[b*3+2] = shz[0] * inv;
    }
}

__inline__ __device__ float warp_sum(float v) {
#pragma unroll
    for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off);
    return v;
}

__global__ void scatter_warp32_sh_kernel(const float* __restrict__ coords,
                                         const float* __restrict__ features,
                                         const float* __restrict__ weight,
                                         const float* __restrict__ bias,
                                         const float* __restrict__ centers,
                                         float* __restrict__ out,
                                         int N, int G, float inv_voxel) {
    int group = blockIdx.x;   // 0..3, eight output channels per block
    int atom  = blockIdx.y;
    int b     = blockIdx.z;
    int tid   = threadIdx.x;
    int lane  = tid & 31;
    int warp  = tid >> 5;
    int fd    = group * 8 + warp;

    __shared__ float xs[32];
    if (tid < 32) xs[tid] = features[(b * N + atom) * 32 + tid];
    __syncthreads();

    const float* c = coords + ((b * N + atom) * 3);
    float halfG = 0.5f * (float)G;
    int ix = (int)((c[0] - centers[b*3+0]) * inv_voxel + halfG);
    int iy = (int)((c[1] - centers[b*3+1]) * inv_voxel + halfG);
    int iz = (int)((c[2] - centers[b*3+2]) * inv_voxel + halfG);
    ix = max(0, min(ix, G - 1));
    iy = max(0, min(iy, G - 1));
    iz = max(0, min(iz, G - 1));

    float sum = warp_sum(xs[lane] * weight[fd * 32 + lane]);
    if (lane == 0) {
        sum += bias[fd];
        int vox = (ix * G + iy) * G + iz;
        int G3 = G * G * G;
        atomicAdd(out + ((b * 32 + fd) * G3 + vox), sum);
    }
}

__global__ void scatter_generic_kernel(const float* __restrict__ coords,
                                       const float* __restrict__ features,
                                       const float* __restrict__ weight,
                                       const float* __restrict__ bias,
                                       const float* __restrict__ centers,
                                       float* __restrict__ out,
                                       int B, int N, int F, int G, float inv_voxel) {
    int atom = blockIdx.x;
    int b = atom / N;
    int n = atom - b * N;
    int d = threadIdx.x;
    const float* c = coords + ((b * N + n) * 3);
    float halfG = 0.5f * (float)G;
    int ix = (int)((c[0] - centers[b*3+0]) * inv_voxel + halfG);
    int iy = (int)((c[1] - centers[b*3+1]) * inv_voxel + halfG);
    int iz = (int)((c[2] - centers[b*3+2]) * inv_voxel + halfG);
    ix = max(0, min(ix, G - 1));
    iy = max(0, min(iy, G - 1));
    iz = max(0, min(iz, G - 1));
    const float* x = features + ((b * N + n) * F);
    int vox = (ix * G + iy) * G + iz;
    int G3 = G * G * G;
    for (int fd = d; fd < F; fd += blockDim.x) {
        float acc = bias[fd];
        const float* w = weight + fd * F;
        for (int j = 0; j < F; ++j) acc += x[j] * w[j];
        atomicAdd(out + ((b * F + fd) * G3 + vox), acc);
    }
}

torch::Tensor voxel_pool_forward(torch::Tensor coords, torch::Tensor features, torch::Tensor weight, torch::Tensor bias, int64_t grid_size, double voxel_size) {
    coords = coords.contiguous();
    features = features.contiguous();
    weight = weight.contiguous();
    bias = bias.contiguous();
    int B = (int)coords.size(0);
    int N = (int)coords.size(1);
    int F = (int)features.size(2);
    int G = (int)grid_size;
    auto out = torch::empty({B, F, G, G, G}, coords.options());
    auto centers = torch::empty({B, 3}, coords.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cudaMemsetAsync(out.data_ptr<float>(), 0, (size_t)B * F * G * G * G * sizeof(float), stream);
    center_kernel<<<B, 256, 0, stream>>>(coords.data_ptr<float>(), centers.data_ptr<float>(), N);
    if (F == 32) {
        dim3 grid(4, N, B);
        scatter_warp32_sh_kernel<<<grid, 256, 0, stream>>>(coords.data_ptr<float>(), features.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), centers.data_ptr<float>(), out.data_ptr<float>(), N, G, (float)(1.0 / voxel_size));
    } else {
        scatter_generic_kernel<<<B * N, (F > 32 ? 128 : 32), 0, stream>>>(coords.data_ptr<float>(), features.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), centers.data_ptr<float>(), out.data_ptr<float>(), B, N, F, G, (float)(1.0 / voxel_size));
    }
    return out;
}
"""

_ext = load_inline(name="voxel_pool_ext_final", cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=["voxel_pool_forward"], extra_cuda_cflags=["-O3", "--use_fast_math"], extra_cflags=["-O3"], 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:
        return _ext.voxel_pool_forward(coords, features, self.linear.weight, self.linear.bias, self.grid_size, float(self.voxel_size))