import os, sys
# sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
import torch
import torch.nn as nn
import torch.nn.functional as F

if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")


class DropBlock(nn.Module):
    """
    this implementation is our implementation for DropBlock described in the DropBlock paper
    You can find the detail implementation and intuition of DropBlock in the following link:
    https://arxiv.org/abs/1810.12890
    """
    def __init__(self, block_size=7, keep_prob=0.9):
        super(DropBlock, self).__init__()
        self.block_size = block_size
        self.keep_prob = keep_prob
        self.gamma = None
        self.kernel_size = (block_size, block_size)
        self.stride = (1, 1)
        self.padding = (block_size // 2, block_size // 2)

    def calculate_gamma(self, x):
        """
        The gamma in DropBlock is the opposite of keep_prob i.e. it controls the number of features to be dropped.
        This function is responsible to adjust the probability based on the keep_prob value since we are dropping
        a continues blocks in feature maps in the DropBlock.
        Args:
            x: feature maps for a minibatch generated by a convolutional module or convolutional layer.

        Returns:
            the gamma which is a float number in [0, 1]
        """
        return (1 - self.keep_prob) * x.shape[-1] ** 2 / \
               (self.block_size ** 2 * (x.shape[-1] - self.block_size + 1) ** 2)

    def forward(self, x):
        if not self.training:
            return x
        if self.gamma is None:
            self.gamma = self.calculate_gamma(x)
        p = torch.ones_like(x) * (self.gamma)
        p = p.to(device)
        # For each feature in the feature map, we will sample from $Bernoulli(\gamma). If the result of this sampling
        # for feature f_{ij} is 0, then keep the feature. If the result of this sampling for f_{ij} is 1,
        # then the entire square region in the mask with the center f_{ij} and the width and height of
        # the square of block_size is set to 0 and therefore dropped.
        m_i_j = torch.bernoulli(p)
        m_i_j = m_i_j.to(device)
        # following line provides the continues blocks that should be dropped with DropBlock denoted as mask here.
        mask = 1 - F.max_pool2d(m_i_j, self.kernel_size, self.stride, self.padding)
        # Normalize the features according to the DorpBlock algorithm described in the DropBlock paper.
        return mask * x * (mask.numel() / mask.sum())
