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 crop_resize_cuda(torch::Tensor input, int out_h, int out_w, int y1, int x1, int h_crop, int w_crop);
        """
        
        cuda_source = """
        #include <cuda_runtime.h>

        __global__ void crop_resize_kernel(
            const float* __restrict__ input,
            float* __restrict__ output,
            int batch, int channels, int in_h, int in_w,
            int out_h, int out_w,
            float scale_y, float scale_x,
            int offset_y, int offset_x
        ) {
            int ow = blockIdx.x * blockDim.x + threadIdx.x;
            int oh = blockIdx.y * blockDim.y + threadIdx.y;
            int bc = blockIdx.z;

            if (ow >= out_w || oh >= out_h || bc >= batch * channels) return;

            float real_y = (oh + 0.5f) * scale_y - 0.5f + offset_y;
            float real_x = (ow + 0.5f) * scale_x - 0.5f + offset_x;

            int y_low = floorf(real_y);
            int x_low = floorf(real_x);
            int y_high = y_low + 1;
            int x_high = x_low + 1;

            float ly = real_y - y_low;
            float lx = real_x - x_low;
            float hy = 1.0f - ly;
            float hx = 1.0f - lx;

            y_low = max(0, min(y_low, in_h-1));
            y_high = max(0, min(y_high, in_h-1));
            x_low = max(0, min(x_low, in_w-1));
            x_high = max(0, min(x_high, in_w-1));

            long long plane_offset = (long long)bc * in_h * in_w;
            float v00 = input[plane_offset + y_low * in_w + x_low];
            float v01 = input[plane_offset + y_low * in_w + x_high];
            float v10 = input[plane_offset + y_high * in_w + x_low];
            float v11 = input[plane_offset + y_high * in_w + x_high];

            float val = v00 * hy * hx + v01 * hy * lx + v10 * ly * hx + v11 * ly * lx;

            long long out_offset = (long long)bc * out_h * out_w + oh * out_w + ow;
            output[out_offset] = val;
        }

        torch::Tensor crop_resize_cuda(torch::Tensor input, int out_h, int out_w, int y1, int x1, int h_crop, int w_crop) {
            int batch = input.size(0);
            int channels = input.size(1);
            int in_h = input.size(2);
            int in_w = input.size(3);

            auto output = torch::empty({batch, channels, out_h, out_w}, input.options());
            
            float scale_y = (float)h_crop / out_h;
            float scale_x = (float)w_crop / out_w;

            dim3 block(16, 16);
            dim3 grid(
                (out_w + 15)/16,
                (out_h + 15)/16,
                batch * channels
            );

            crop_resize_kernel<<<grid, block>>>(
                input.data_ptr<float>(),
                output.data_ptr<float>(),
                batch, channels, in_h, in_w, out_h, out_w,
                scale_y, scale_x, y1, x1
            );

            return output;
        }
        """

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not x.is_contiguous(): x = x.contiguous()
        return self.op.crop_resize_cuda(x, 128, 128, 100, 100, 300, 300)