import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor gridsample1d_cuda(torch::Tensor input, torch::Tensor grid);
        """
        
        cuda_source = """
        #include <cuda_runtime.h>

        #define BLOCK_SIZE 256

        __device__ __forceinline__ float get_pixel_1d(
            const float* __restrict__ data, 
            int w, 
            int W, 
            long long channel_offset
        ) {
            if (w >= 0 && w < W) {
                return data[channel_offset + w];
            }
            return 0.0f; // Padding zeros
        }

      
        __global__ void gridsample1d_stable_kernel(
            const float* __restrict__ input,
            const float* __restrict__ grid,
            float* __restrict__ output,
            int total_spatial, // B * W_out
            int batch,
            int channels,
            int w_in,
            int w_out
        ) {
            int idx = blockIdx.x * blockDim.x + threadIdx.x;
            if (idx >= total_spatial) return;

            int w = idx % w_out;
            int b = idx / w_out;

          
            long long grid_idx = (long long)b * w_out + w;
            float g_val = grid[grid_idx];

            double x = (double)g_val;
            
            // align_corners=True: ((x + 1) * (W - 1)) / 2
            double real_x = (x + 1.0) * (w_in - 1.0) * 0.5;

            int x_w = floor(real_x + 1e-5);
            int x_e = x_w + 1;


            double d_e = real_x - x_w; // distance to west (weight for east)
            double w_w = 1.0 - d_e;    // weight for west
            double w_e = d_e;          // weight for east

            long long in_batch_base = (long long)b * (channels * w_in);
            long long out_batch_base = (long long)b * (channels * w_out);

  
            for (int c = 0; c < channels; ++c) {
                long long in_c_offset = in_batch_base + c * w_in;
                
             
                float v_w = get_pixel_1d(input, x_w, w_in, in_c_offset);
                float v_e = get_pixel_1d(input, x_e, w_in, in_c_offset);

             
                double val = (double)v_w * w_w + (double)v_e * w_e;

                long long out_idx = out_batch_base + c * w_out + w;
                output[out_idx] = (float)val;
            }
        }

        torch::Tensor gridsample1d_cuda(torch::Tensor input, torch::Tensor grid) {
            int batch = input.size(0);
            int channels = input.size(1);
            int w_in = input.size(2);
            int w_out = grid.size(1); // Grid [N, W_out, 1]
            
            auto output = torch::empty({batch, channels, w_out}, input.options());
            
            long long total_spatial = (long long)batch * w_out;
            
            const int block_size = 256;
            const int grid_size = (total_spatial + block_size - 1) / block_size;

            gridsample1d_stable_kernel<<<grid_size, block_size>>>(
                input.data_ptr<float>(),
                grid.data_ptr<float>(),
                output.data_ptr<float>(),
                total_spatial,
                batch, channels, w_in, w_out
            );
            
            return output;
        }
        """

        self.op = load_inline(
            name="gridsample1d_stable_v1",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["gridsample1d_cuda"],
            extra_cuda_cflags=["-O3"], 
            verbose=False
        )

    def forward(self, x: torch.Tensor, grid: torch.Tensor) -> torch.Tensor:
        if not x.is_contiguous(): x = x.contiguous()
        if not grid.is_contiguous(): grid = grid.contiguous()
        return self.op.gridsample1d_cuda(x, grid)