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 * \
               x.shape[-1] / (x.shape[-1] - self.block_size + 1)

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor dropblock1d_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 dropblock1d_kernel(
            const float* __restrict__ input,
            const float* __restrict__ rand,
            float* __restrict__ output,
            long long total_elements,
            int channels,
            int length,
            int block_radius,
            float gamma,
            float scale
        ) {
            long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
            if (idx >= total_elements) return;

         
            int l = idx % length;
            int c = (idx / length) % channels;
            // Batch not needed for this layout

        
            int r = block_radius;
            int l_min = max(0, l - r);
            int l_max = min(length - 1, l + r);
            
            long long plane_offset = (idx / length) * length;
            
            bool dropped = false;
            
            for (int i = l_min; i <= l_max; ++i) {
          
                if (rand[plane_offset + i] < gamma) {
                    dropped = true;
                    break;
                }
            }

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

        torch::Tensor dropblock1d_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 length = input.size(2);
            
            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;

            dropblock1d_kernel<<<grid_size, threads>>>(
                input.data_ptr<float>(),
                rand.data_ptr<float>(),
                output.data_ptr<float>(),
                total_elements,
                channels,
                length,
                radius, gamma, scale
            );
            
            return output;
        }
        """

        self.op = load_inline(
            name="dropblock1d_fused_v1",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["dropblock1d_cuda"],
            extra_cuda_cflags=["-O3"],
            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.dropblock1d_cuda(x, rand, self.block_size, self.gamma, scale)