##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

## Created by: Hang Zhang

## Email: zhanghang0704@gmail.com

## Copyright (c) 2020

##

## LICENSE file in the root directory of this source tree 

##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

import mxnet as mx

from functools import partial

from mxnet.gluon.nn import MaxPool1D, Block, HybridBlock

__all__ = ['DropBlock', 'set_drop_prob', 'DropBlockScheduler']

class DropBlock(HybridBlock):

    def __init__(self, drop_prob, block_size, c, h, w):

        super().__init__()

        self.drop_prob = drop_prob

        self.block_size = block_size

        self.c, self.h, self.w = c, h, w

        self.numel = c * h * w

        pad_h = max((block_size - 1), 0)

        pad_w = max((block_size - 1), 0)

        self.padding = (pad_h//2, pad_h-pad_h//2, pad_w//2, pad_w-pad_w//2)

        self.dtype = 'float32'

    def hybrid_forward(self, F, x):

        if not mx.autograd.is_training() or self.drop_prob <= 0:

            return x

        gamma = self.drop_prob * (self.h * self.w) / (self.block_size ** 2) / \

            ((self.w - self.block_size + 1) * (self.h - self.block_size + 1))

        # generate mask

        mask = F.random.uniform(0, 1, shape=(1, self.c, self.h, self.w), dtype=self.dtype) < gamma

        mask = F.Pooling(mask, pool_type='max',

                         kernel=(self.block_size, self.block_size), pad=self.padding)

        mask = 1 - mask

        y = F.broadcast_mul(F.broadcast_mul(x, mask),

                            (1.0 * self.numel / mask.sum(axis=0, exclude=True).expand_dims(1).expand_dims(1).expand_dims(1)))

        return y

    def cast(self, dtype):

        super(DropBlock, self).cast(dtype)

        self.dtype = dtype

    def __repr__(self):

        reprstr = self.__class__.__name__ + '(' + \

            'drop_prob: {}, block_size{}'.format(self.drop_prob, self.block_size) +')'

        return reprstr

def set_drop_prob(drop_prob, module):

    """

    Example:

        from functools import partial

        apply_drop_prob = partial(set_drop_prob, 0.1)

        net.apply(apply_drop_prob)

    """

    if isinstance(module, DropBlock):

        module.drop_prob = drop_prob

class DropBlockScheduler(object):

    def __init__(self, net, start_prob, end_prob, num_epochs):

        self.net = net

        self.start_prob = start_prob

        self.end_prob = end_prob

        self.num_epochs = num_epochs

    def __call__(self, epoch):

        ratio = self.start_prob + 1.0 * (self.end_prob - self.start_prob) * (epoch + 1) / self.num_epochs

        assert (ratio >= 0 and ratio <= 1)

        apply_drop_prob = partial(set_drop_prob, ratio)

        self.net.apply(apply_drop_prob)

        self.net.hybridize()