import torch
from torch.utils.cpp_extension import load_inline

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

__global__ void pointpillars_scatter_kernel_opt(
    const float* __restrict__ feats,
    const int* __restrict__ coords,
    float* __restrict__ output,
    int M, int C, int H, int W, 
    int out_stride_c, int out_stride_b) 
{
    int m = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (m < M) {
     
        const int4* coord_ptr = reinterpret_cast<const int4*>(coords);
        int4 c_val = coord_ptr[m];
        
        int b = c_val.x;
        // int z = c_val.y; 
        int y = c_val.z;
        int x = c_val.w;
        
        if (b >= 0 && y >= 0 && y < H && x >= 0 && x < W) {
            long out_base_idx = (long)b * out_stride_b + (long)y * W + x;
            
            // Input feature pointer
            const float* feat_ptr = feats + m * C;
      
            int vec_loops = C / 4;
            const float4* feat_ptr_4 = reinterpret_cast<const float4*>(feat_ptr);
            
            for (int k = 0; k < vec_loops; ++k) {
                float4 val = feat_ptr_4[k];
                
          
                int c_base = k * 4;

                output[out_base_idx + (long)(c_base + 0) * out_stride_c] = val.x;
                output[out_base_idx + (long)(c_base + 1) * out_stride_c] = val.y;
                output[out_base_idx + (long)(c_base + 2) * out_stride_c] = val.z;
                output[out_base_idx + (long)(c_base + 3) * out_stride_c] = val.w;
            }
            
            for (int k = vec_loops * 4; k < C; ++k) {
                output[out_base_idx + (long)k * out_stride_c] = feat_ptr[k];
            }
        }
    }
}

torch::Tensor pointpillars_scatter_cuda(torch::Tensor feats, torch::Tensor coords, int B, int H, int W) {
    int M = feats.size(0);
    int C = feats.size(1);

    auto output = torch::zeros({B, C, H, W}, feats.options());
    

    int out_stride_c = H * W;     // Stride between channels
    int out_stride_b = C * H * W; // Stride between batches
    
    // Config: 1 Thread per Pillar
    const int block = 256;
    const int grid = (M + block - 1) / block;
    
    pointpillars_scatter_kernel_opt<<<grid, block>>>(
        feats.data_ptr<float>(),
        coords.data_ptr<int>(),
        output.data_ptr<float>(),
        M, C, H, W,
        out_stride_c, out_stride_b
    );
    
    return output;
}
"""

cpp_source = "torch::Tensor pointpillars_scatter_cuda(torch::Tensor feats, torch::Tensor coords, int B, int H, int W);"

module = load_inline(
    name="pointpillars_scatter_extension_v3",
    cpp_sources=cpp_source,
    cuda_sources=source,
    functions=["pointpillars_scatter_cuda"],
    verbose=True,
    with_cuda=True
)

class ModelNew(torch.nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.H = 512
        self.W = 512
        self.B = 4
        self.op = module

    def forward(self, feats, coords):
        return self.op.pointpillars_scatter_cuda(
            feats.contiguous(), 
            coords.int().contiguous(), 
            self.B, self.H, self.W
        )