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

    def calculate_gamma(self, x):
        D, H, W = x.shape[-3], x.shape[-2], x.shape[-1]
        vol = (D - self.block_size + 1) * (H - self.block_size + 1) * (W - self.block_size + 1)
        if vol <= 0: return 0.0
        return (1.0 - self.keep_prob) / (self.block_size ** 3) * (D * H * W) / vol

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

        // Block Output Size: 4x8x8 = 256 threads
        #define TILE_D 4
        #define TILE_H 8
        #define TILE_W 8
        
        // Kernel Radius (BlockSize=5 -> Radius=2)
        #define RADIUS 2
        
        // Shared Memory Size for Rand
        #define SMEM_D (TILE_D + RADIUS * 2) // 4+4=8
        #define SMEM_H (TILE_H + RADIUS * 2) // 8+4=12
        #define SMEM_W (TILE_W + RADIUS * 2) // 8+4=12

        __global__ void dropblock3d_smem_kernel(
            const float* __restrict__ input,
            const float* __restrict__ rand,
            float* __restrict__ output,
            int batch, int channels, int depth, int height, int width,
            float gamma,
            float scale
        ) {
            // Grid Mapping
            // Block: 4x8x8
            int tx = threadIdx.x;
            int ty = threadIdx.y;
            int tz = threadIdx.z;

            // Grid.x -> W blocks, Grid.y -> H blocks, Grid.z -> D blocks * B * C
            int w_blocks = (width + TILE_W - 1) / TILE_W;
            int h_blocks = (height + TILE_H - 1) / TILE_H;
            
            int bz = blockIdx.z;
            int by = blockIdx.y;
            int bx = blockIdx.x;
            
            int tmp = bz;
            int c = tmp % channels; tmp /= channels;
            int b = tmp;

            int od_base = by / h_blocks * TILE_D;
            int oh_base = by % h_blocks * TILE_H;
            int ow_base = bx * TILE_W;
            
            int od = od_base + tz;
            int oh = oh_base + ty;
            int ow = ow_base + tx;

            // Shared Memory for Rand Tile
            __shared__ float smem[SMEM_D][SMEM_H][SMEM_W];

            // 1. Cooperative Loading (Rand -> Smem)
            int rand_d_start = od_base - RADIUS;
            int rand_h_start = oh_base - RADIUS;
            int rand_w_start = ow_base - RADIUS;
            
            int tid = tz * (TILE_H * TILE_W) + ty * TILE_W + tx; // 0..255
            int num_smem_elements = SMEM_D * SMEM_H * SMEM_W;
            
            long long rand_plane_offset = (long long)b * (channels * depth * height * width) + c * (depth * height * width);

            for (int i = tid; i < num_smem_elements; i += 256) {
                int sz = i / (SMEM_H * SMEM_W);
                int rem_s = i % (SMEM_H * SMEM_W);
                int sy = rem_s / SMEM_W;
                int sx = rem_s % SMEM_W;
                
                int gd = rand_d_start + sz;
                int gh = rand_h_start + sy;
                int gw = rand_w_start + sx;
                
                float val = 1.0f; // Default > gamma
                if (gd >= 0 && gd < depth && gh >= 0 && gh < height && gw >= 0 && gw < width) {
                    val = rand[rand_plane_offset + gd * (height * width) + gh * width + gw];
                }
                smem[sz][sy][sx] = val;
            }

            __syncthreads();

            // 2. Compute
            if (od < depth && oh < height && ow < width) {
                
                bool dropped = false;
                
                // Window Search in Smem
                // Smem local coord for center: tz+R, ty+R, tx+R
                
                // We check a 5x5x5 window. The window start in Smem: tz, ty, tx
                #pragma unroll
                for (int z = tz; z < tz + 2*RADIUS + 1; ++z) {
                    #pragma unroll
                    for (int y = ty; y < ty + 2*RADIUS + 1; ++y) {
                        #pragma unroll
                        for (int x = tx; x < tx + 2*RADIUS + 1; ++x) {
                            if (smem[z][y][x] < gamma) {
                                dropped = true;
                                goto end_check;
                            }
                        }
                    }
                }
                end_check:;

                // 3. Apply & Write
                long long io_idx = rand_plane_offset + od * (height * width) + oh * width + ow;
                if (dropped) {
                    output[io_idx] = 0.0f;
                } else {
                    output[io_idx] = input[io_idx] * scale;
                }
            }
        }

        torch::Tensor dropblock3d_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 depth = input.size(2);
            int height = input.size(3);
            int width = input.size(4);
            
            auto output = torch::empty_like(input);
            
            dim3 block(TILE_W, TILE_H, TILE_D);
            
            // Grid Mapping
            // X: Width Blocks
            // Y: Height * Depth Blocks (Flattened)
            // Z: Batch * Channel Blocks (Flattened)
            dim3 grid(
                (width + TILE_W - 1) / TILE_W,
                ((height + TILE_H - 1) / TILE_H) * ((depth + TILE_D - 1) / TILE_D),
                batch * channels
            );

            dropblock3d_smem_kernel<<<grid, block>>>(
                input.data_ptr<float>(),
                rand.data_ptr<float>(),
                output.data_ptr<float>(),
                batch, channels, depth, height, width,
                gamma, scale
            );
            
            return output;
        }
        """

        self.op = load_inline(
            name="dropblock3d_smem_v1",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["dropblock3d_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.dropblock3d_cuda(x, rand, self.block_size, self.gamma, scale)