import torch
from torch.utils.cpp_extension import load_inline

ball_query_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void ball_query_kernel(
    int B, int N, int M, float radius_sq, int nsample,
    const float* __restrict__ xyz,     // [B, N, 3]
    const float* __restrict__ new_xyz, // [B, M, 3]
    int* __restrict__ idx_out)         // [B, M, nsample]
{
    // Grid X: M (Query Points)
    // Grid Y: B (Batch)
    int m_idx = blockIdx.x * blockDim.x + threadIdx.x;
    int b_idx = blockIdx.y;
    
    if (m_idx >= M || b_idx >= B) return;
    
    // Base pointers
    const float* new_xyz_ptr = new_xyz + b_idx * M * 3 + m_idx * 3;
    const float* xyz_ptr     = xyz     + b_idx * N * 3;
    int* out_ptr             = idx_out + b_idx * M * nsample + m_idx * nsample;
    
    // 1. Load Query Point to Registers
    float qx = new_xyz_ptr[0];
    float qy = new_xyz_ptr[1];
    float qz = new_xyz_ptr[2];
    
    // 2. Iterate over all N points
    int cnt = 0;
    
    for (int k = 0; k < N; ++k) {
        float x = xyz_ptr[k * 3 + 0];
        float y = xyz_ptr[k * 3 + 1];
        float z = xyz_ptr[k * 3 + 2];
        
        float d2 = (qx - x) * (qx - x) + 
                   (qy - y) * (qy - y) + 
                   (qz - z) * (qz - z);
                   
        if (d2 < radius_sq) {
            if (cnt == 0) {
                // Pre-fill the buffer with the first found point
                // This handles the "duplicate first point" logic automatically
                // We will fill the rest as we find them
                for (int l = 0; l < nsample; ++l) {
                    out_ptr[l] = k;
                }
            }
            
            // Record current point
            out_ptr[cnt] = k;
            cnt++;
            
            // Early exit if full
            if (cnt >= nsample) break;
        }
    }
    
    // If cnt == 0, we didn't find any point. 
    // Standard behavior is usually index 0.
    if (cnt == 0) {
        for (int l = 0; l < nsample; ++l) {
            out_ptr[l] = 0;
        }
    }
}

torch::Tensor ball_query_cuda(torch::Tensor xyz, torch::Tensor new_xyz, float radius, int nsample) {
    int B = xyz.size(0);
    int N = xyz.size(1);
    int M = new_xyz.size(1);
    
    // Output indices
    auto idx = torch::zeros({B, M, nsample}, torch::dtype(torch::kInt32).device(xyz.device()));
    
    float radius_sq = radius * radius;
    
    // Config
    // Block: 256 threads
    // Grid X: M / 256
    // Grid Y: B
    const int block = 256;
    dim3 grid((M + block - 1) / block, B);
    
    ball_query_kernel<<<grid, block>>>(
        B, N, M, radius_sq, nsample,
        xyz.data_ptr<float>(),
        new_xyz.data_ptr<float>(),
        idx.data_ptr<int>()
    );
    
    return idx;
}
"""

cpp_source = "torch::Tensor ball_query_cuda(torch::Tensor xyz, torch::Tensor new_xyz, float radius, int nsample);"

ball_query_module = load_inline(
    name="ball_query_extension",
    cpp_sources=cpp_source,
    cuda_sources=ball_query_source,
    functions=["ball_query_cuda"],
    verbose=True,
    with_cuda=True
)

class ModelNew(torch.nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.radius = 20.0
        self.nsample = 32
        self.cuda_op = ball_query_module

    def forward(self, xyz, new_xyz):
        return self.cuda_op.ball_query_cuda(
            xyz.contiguous(), 
            new_xyz.contiguous(), 
            self.radius, 
            self.nsample
        )