import os
os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;7.0;7.5;8.0;8.6+PTX"

import torch
from torch.utils.cpp_extension import load_inline

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

__global__ void bilinear_interpolate_kernel(
    const torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> input, // (N, C, H_in, W_in)
    const torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> grid,  // (N, H_out, W_out, 2)
    torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> output,      // (N, C, H_out, W_out)
    const int C, const int H_out, const int W_out,
    const int H_in, const int W_in
) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total_elements = grid.size(0) * C * H_out * W_out; // N * C * H_out * W_out

    if (idx >= total_elements) return;

    const int n = idx / (C * H_out * W_out);
    const int rem1 = idx % (C * H_out * W_out);
    const int c = rem1 / (H_out * W_out);
    const int rem2 = rem1 % (H_out * W_out);
    const int h_out = rem2 / W_out;
    const int w_out = rem2 % W_out;

    const float x_f = grid[n][h_out][w_out][0]; // in [-1, 1]
    const float y_f = grid[n][h_out][w_out][1]; // in [-1, 1]

    const float x = (x_f + 1.0f) * 0.5f * (W_in - 1);
    const float y = (y_f + 1.0f) * 0.5f * (H_in - 1);

    const int x0 = static_cast<int>(x);
    const int y0 = static_cast<int>(y);
    const int x1 = x0 + 1;
    const int y1 = y0 + 1;

    const float wa = (x1 - x) * (y1 - y);
    const float wb = (x - x0) * (y1 - y);
    const float wc = (x1 - x) * (y - y0);
    const float wd = (x - x0) * (y - y0);

    float value = 0.0f;
    if (x0 >= 0 && x0 < W_in && y0 >= 0 && y0 < H_in) {
        value += wa * input[n][c][y0][x0];
    }
    if (x1 >= 0 && x1 < W_in && y0 >= 0 && y0 < H_in) {
        value += wb * input[n][c][y0][x1];
    }
    if (x0 >= 0 && x0 < W_in && y1 >= 0 && y1 < H_in) {
        value += wc * input[n][c][y1][x0];
    }
    if (x1 >= 0 && x1 < W_in && y1 >= 0 && y1 < H_in) {
        value += wd * input[n][c][y1][x1];
    }

    // 8. 写入输出
    output[n][c][h_out][w_out] = value;
}

// C++包装函数
torch::Tensor bilinear_interpolate_cuda(torch::Tensor input, torch::Tensor grid) {
    TORCH_CHECK(input.is_cuda(), "Input image must be a CUDA tensor");
    TORCH_CHECK(grid.is_cuda(), "Grid must be a CUDA tensor");
    TORCH_CHECK(input.dim() == 4, "Input image must be a 4D tensor");
    TORCH_CHECK(grid.dim() == 4, "Grid must be a 4D tensor");

    const int N = input.size(0);
    const int C = input.size(1);
    const int H_in = input.size(2);
    const int W_in = input.size(3);
    const int H_out = grid.size(1);
    const int W_out = grid.size(2);

    auto output = torch::zeros({N, C, H_out, W_out}, input.options());

    const int threads = 256;
    const int total_elements = N * C * H_out * W_out;
    const int blocks = (total_elements + threads - 1) / threads;

    bilinear_interpolate_kernel<<<blocks, threads>>>(
        input.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
        grid.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
        output.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
        C, H_out, W_out, H_in, W_in
    );

    return output;
}
"""

bilinear_cpp_source = """
torch::Tensor bilinear_interpolate_cuda(torch::Tensor input, torch::Tensor grid);
"""

# 编译模块
bilinear_module = load_inline(
    name="bilinear_interpolate",
    cpp_sources=bilinear_cpp_source,
    cuda_sources=bilinear_source,
    functions=["bilinear_interpolate_cuda"],
    verbose=True,
    extra_cuda_cflags=["-O3"]
)

class ModelNew(torch.nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.bilinear = bilinear_module

    def forward(self, input_image, grid):
        # 确保输入是连续的
        if not input_image.is_contiguous(): input_image = input_image.contiguous()
        if not grid.is_contiguous(): grid = grid.contiguous()
        
        return self.bilinear.bilinear_interpolate_cuda(input_image, grid)
