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

cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void voxel_pool_kernel(
    const float* __restrict__ coords,
    const float* __restrict__ features,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ voxel_grid,
    int B, int N, int F, int G, float voxel_size, float inv_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;

    // Compute center for batch b
    float cx = 0.0f, cy = 0.0f, cz = 0.0f;
    for (int i = 0; i < N; ++i) {
        int base = (b * N + i) * 3;
        cx += coords[base];
        cy += coords[base + 1];
        cz += coords[base + 2];
    }
    cx /= N; cy /= N; cz /= N;

    int base_c = (b * N + n) * 3;
    float x = (coords[base_c]     - cx) * inv_voxel_size + G * 0.5f;
    float y = (coords[base_c + 1] - cy) * inv_voxel_size + G * 0.5f;
    float z = (coords[base_c + 2] - cz) * inv_voxel_size + G * 0.5f;

    int ix = (int)floorf(x);
    int iy = (int)floorf(y);
    int iz = (int)floorf(z);
    if (ix < 0) ix = 0; if (ix >= G) ix = G-1;
    if (iy < 0) iy = 0; if (iy >= G) iy = G-1;
    if (iz < 0) iz = 0; if (iz >= G) iz = G-1;

    // Linear projection for this point's features
    const float* feat_ptr = features + (b * N + n) * F;
    float* out_ptr_base = voxel_grid + b * F * G * G * G;

    for (int d = 0; d < F; ++d) {
        float val = bias[d];
        const float* w_row = weight + d * F;
        for (int k = 0; k < F; ++k) {
            val += feat_ptr[k] * w_row[k];
        }
        float* voxel_addr = out_ptr_base + d * G*G*G + ix * G*G + iy * G + iz;
        atomicAdd(voxel_addr, val);
    }
}

torch::Tensor voxel_pool_cuda(
    torch::Tensor coords,
    torch::Tensor features,
    torch::Tensor weight,
    torch::Tensor bias,
    int grid_size,
    float voxel_size
) {
    int B = coords.size(0);
    int N = coords.size(1);
    int F = features.size(2);
    int G = grid_size;
    float inv_voxel_size = 1.0f / voxel_size;

    auto voxel_grid = torch::zeros({B, F, G, G, G}, coords.options());

    int total = B * N;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;

    voxel_pool_kernel<<<blocks, threads>>>(
        coords.data_ptr<float>(),
        features.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        voxel_grid.data_ptr<float>(),
        B, N, F, G, voxel_size, inv_voxel_size
    );

    return voxel_grid;
}
'''

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

mod = load_inline(
    name='voxel_pool_mod',
    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:
        return mod.voxel_pool_cuda(
            coords, features,
            self.linear.weight, self.linear.bias,
            self.grid_size, self.voxel_size
        )
```