import torch
import torch.nn as nn
import torch.nn.functional as F


BATCH = 8
CHANNELS = 32
DEPTH = 32
HEIGHT = 64
WIDTH = 64
BLOCK_SIZE = 5
KEEP_PROB = 0.9

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.block_size = BLOCK_SIZE
        self.keep_prob = KEEP_PROB
        self.gamma = None

    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 forward(self, x: torch.Tensor, rand_tensor: torch.Tensor) -> torch.Tensor:
        if self.gamma is None:
            self.gamma = self.calculate_gamma(x)

        mask = (rand_tensor < self.gamma).float()

        pad = self.block_size // 2
        mask = F.max_pool3d(mask, kernel_size=self.block_size, stride=1, padding=pad)
   
        mask = 1.0 - mask
     
        scale = 1.0 / self.keep_prob
        
        return x * mask * scale

def get_inputs():
    x = torch.randn(BATCH, CHANNELS, DEPTH, HEIGHT, WIDTH, device='cuda', dtype=torch.float32)
    rand = torch.rand_like(x)
    return [x, rand]

def get_init_inputs():
    return []