import torch
from torch.utils.cpp_extension import load_inline

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



__global__ void gather_nd_float4_kernel(
    const float* __restrict__ params, 
    const int* __restrict__ indices, 
    float* __restrict__ output, 
    int N, 
    int C,
    int stride_b, int stride_h, int stride_w) 
{

    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < N) {

        int base_idx = idx * 3;
        int b = indices[base_idx];
        int h = indices[base_idx + 1];
        int w = indices[base_idx + 2];

        long long src_offset = (long long)b * stride_b + (long long)h * stride_h + (long long)w * stride_w;
        

        long long dst_offset = (long long)idx * C;

    
        const float4* src_ptr = reinterpret_cast<const float4*>(params + src_offset);
        float4* dst_ptr = reinterpret_cast<float4*>(output + dst_offset);

        int vec_len = C / 4;
   
        for (int i = 0; i < vec_len; ++i) {
            dst_ptr[i] = src_ptr[i];
        }
    }
}

torch::Tensor gather_nd_cuda(torch::Tensor params, torch::Tensor indices) {
    
    int N = indices.size(0);
    int C = params.size(3);
    int H = params.size(1);
    int W = params.size(2);
    

    int stride_w = C;
    int stride_h = W * C;
    int stride_b = H * W * C;

    auto output = torch::empty({N, C}, params.options());
    
    const int block = 256;
    const int grid = (N + block - 1) / block;
 
    gather_nd_float4_kernel<<<grid, block>>>(
        params.data_ptr<float>(),
        indices.data_ptr<int>(),
        output.data_ptr<float>(),
        N, C,
        stride_b, stride_h, stride_w
    );
    
    return output;
}
"""

cpp_source = "torch::Tensor gather_nd_cuda(torch::Tensor params, torch::Tensor indices);"

gather_nd_module = load_inline(
    name="gather_nd_extension",
    cpp_sources=cpp_source,
    cuda_sources=gather_nd_source,
    functions=["gather_nd_cuda"],
    verbose=True,
    with_cuda=True
)

class ModelNew(torch.nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.cuda_op = gather_nd_module

    def forward(self, params, indices):

        return self.cuda_op.gather_nd_cuda(params.contiguous(), indices.int().contiguous())