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.block_size = 7
        self.keep_prob = 0.9
        self.gamma = None
        self._compile_cuda_kernel()

    def calculate_gamma(self, x):
        return (1.0 - self.keep_prob) / (self.block_size ** 2) * \
               (x.shape[-2] * x.shape[-1]) / \
               ((x.shape[-2] - self.block_size + 1) * (x.shape[-1] - self.block_size + 1))

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor dropblock2d_cuda(torch::Tensor input, torch::Tensor rand, int block_size, float gamma, float scale);
        """
        
        cuda_source = """
        #include <cuda_runtime.h>

        #define BLOCK_SIZE 256

        __global__ void dropblock2d_pixel_kernel(
            const float* __restrict__ input,
            const float* __restrict__ rand,
            float* __restrict__ output,
            long long total_elements,
            int height, 
            int width,
            int block_radius,
            float gamma,
            float scale
        ) {
            long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
            if (idx >= total_elements) return;

            int w = idx % width;
            int h = (idx / width) % height;

            int r = block_radius;
            int y_min = max(0, h - r);
            int y_max = min(height - 1, h + r);
            int x_min = max(0, w - r);
            int x_max = min(width - 1, w + r);
  
            long long plane_idx = idx / (height * width);
            long long plane_offset = plane_idx * (height * width);
            
            bool dropped = false;
            
            for (int y = y_min; y <= y_max; ++y) {
                for (int x = x_min; x <= x_max; ++x) {
                    if (rand[plane_offset + y * width + x] < gamma) {
                        dropped = true;
                        goto end_check;
                    }
                }
            }
            end_check:;

            if (dropped) {
                output[idx] = 0.0f;
            } else {
                output[idx] = input[idx] * scale;
            }
        }

        torch::Tensor dropblock2d_cuda(torch::Tensor input, torch::Tensor rand, int block_size, float gamma, float scale) {
            int batch = input.size(0);
            int channels = input.size(1);
            int height = input.size(2);
            int width = input.size(3);
            
            auto output = torch::empty_like(input);
            
            long long total_elements = input.numel();
            
            const int threads = 256;

            long long grid_size_ll = (total_elements + threads - 1) / threads;
        
            int grid_size = (int)grid_size_ll;

            int radius = block_size / 2;

            dropblock2d_pixel_kernel<<<grid_size, threads>>>(
                input.data_ptr<float>(),
                rand.data_ptr<float>(),
                output.data_ptr<float>(),
                total_elements,
                height, width,
                radius, gamma, scale
            );
            
            return output;
        }
        """

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

    def forward(self, x: torch.Tensor, rand: torch.Tensor) -> torch.Tensor:
        if self.gamma is None:
            self.gamma = self.calculate_gamma(x)
        if not x.is_contiguous(): x = x.contiguous()
        if not rand.is_contiguous(): rand = rand.contiguous()
        
        scale = 1.0 / self.keep_prob
        
        return self.op.dropblock2d_cuda(x, rand, self.block_size, self.gamma, scale)