""" EfficientNet, MobileNetV3, etc Blocks

Hacked together by / Copyright 2019, Ross Wightman
"""
import random

import torch
import torch.nn as nn
from torch.nn import functional as F
from ..build import BLOCKS, PARAMS_MAPPING, ABBREVS


def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by_keep: bool = True):
    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
    'survival rate' as the argument.

    """
    if drop_prob == 0. or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
    if keep_prob > 0.0 and scale_by_keep:
        random_tensor.div_(keep_prob)
    return x * random_tensor


def make_divisible(v, divisor=8, min_value=None, round_limit=.9):
    min_value = min_value or divisor
    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
    # Make sure that round down does not go down by more than 10%.
    if new_v < round_limit * v:
        new_v += divisor
    return new_v

@BLOCKS.register
@PARAMS_MAPPING.register
@ABBREVS.register('ds')
class DepthwiseSeparableConv(nn.Module):
    """ DepthwiseSeparable block
    Used for DS convs in MobileNet-V1 and in the place of IR blocks that have no expansion
    (factor of 1.0). This is an alternative to having a IR with an optional first pw conv.
    """
    _PARAMS_DICT = {
        'i': 'in_channels',
        'o': 'out_channels',
        'k': 'dw_kernel_size',
        's': 'stride',
        'noskip': 'noskip',
        'p': 'pw_kernel_size',
        'g': 'pw_group_num',
        'n': 'act_layer',
        'd': 'dilation',
        'c': 'dw_conv'
    }

    def __init__(self, in_channels, out_channels, dw_conv='Conv', dw_kernel_size=3, stride=1, dilation=1, pad_type='',
                 noskip=False, pw_kernel_size=1, pw_group_num=1, pw_act=False, act_layer='ReLU',
                 norm_layer=nn.BatchNorm2d, se_layer=None, drop_path_rate=0., **kwargs):
        super(DepthwiseSeparableConv, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.has_residual = (stride == 1 and in_channels == out_channels) and not noskip
        self.has_pw_act = pw_act  # activation after point-wise conv
        self.drop_path_rate = drop_path_rate
        if dw_conv == 'Conv':
            dwconv = nn.Conv2d

        padding = ((stride - 1) + dilation * (dw_kernel_size - 1)) // 2

        if norm_layer is not None:
            self.conv_dw = dwconv(in_channels, in_channels, dw_kernel_size, stride=stride,
                                        dilation=dilation, padding=padding, groups=in_channels, bias=False)
            self.bn1 = nn.BatchNorm2d(in_channels)
            self.act1 = nn.ReLU()

            self.conv_pw = nn.Conv2d(in_channels, out_channels, pw_kernel_size, groups=pw_group_num, bias=False)
            self.bn2 = nn.BatchNorm2d(out_channels)
            self.layers = ['conv_dw', 'bn1', 'act1', 'conv_pw', 'bn2']
        else:
            self.conv_dw = dwconv(in_channels, in_channels, dw_kernel_size, stride=stride,
                                        dilation=dilation, padding=padding, groups=in_channels, bias=True)
            self.act1 = nn.ReLU()

            self.conv_pw = nn.Conv2d(in_channels, out_channels, pw_kernel_size, groups=pw_group_num, bias=True)
            self.layers = ['conv_dw', 'act1', 'conv_pw']
        if self.has_pw_act:
            self.act2 = nn.ReLU()
            self.layers.append('act2')

    def feature_info(self, location):
        if location == 'expansion':  # after SE, input to PW
            info = dict(module='conv_pw', hook_type='forward_pre', num_chs=self.in_channels)
        else:  # location == 'bottleneck', block output
            info = dict(module='', hook_type='', num_chs=self.out_channels)
        return info

    def forward(self, x):
        shortcut = x

        for name in self.layers:
            x = getattr(self, name)(x)

        if self.has_residual:
            if self.drop_path_rate > 0.:
                x = drop_path(x, self.drop_path_rate, self.training)
            x += shortcut
        return x

@BLOCKS.register
@PARAMS_MAPPING.register
@ABBREVS.register('ir')
class InvertedResidual(nn.Module):
    """ Inverted residual block w/ optional SE

    Originally used in MobileNet-V2 - https://arxiv.org/abs/1801.04381v4, this layer is often
    referred to as 'MBConv' for (Mobile inverted bottleneck conv) and is also used in
      * MNasNet - https://arxiv.org/abs/1807.11626
      * EfficientNet - https://arxiv.org/abs/1905.11946
      * MobileNet-V3 - https://arxiv.org/abs/1905.02244
    """
    _PARAMS_DICT = {
        'i': 'in_channels',
        'o': 'out_channels',
        'k': 'dw_kernel_size',
        's': 'stride',
        'noskip': 'noskip',
        'e': 'exp_ratio',
        'a': 'exp_kernel_size',
        'p': 'pw_kernel_size',
        'g': 'pw_group_num',
        'n': 'act_layer',
        'd': 'dilation',
        'c': 'dw_conv',
    }

    def __init__(
            self, in_channels, out_channels, dw_conv='Conv', dw_kernel_size=3, stride=1, dilation=1, pad_type='',
            noskip=False, exp_ratio=1.0, exp_kernel_size=1, pw_kernel_size=1, pw_group_num=1, act_layer='ReLU',
            norm_layer=nn.BatchNorm2d, se_layer=None, conv_kwargs=None, drop_path_rate=0., **kwargs):
        super(InvertedResidual, self).__init__()
        conv_kwargs = conv_kwargs or {}
        self.has_residual = (in_channels == out_channels and stride == 1) and not noskip
        self.drop_path_rate = drop_path_rate
        self.in_channels = in_channels
        self.out_channels = out_channels
        if dw_conv == 'Conv':
            dwconv = nn.Conv2d

        mid_chs = make_divisible(in_channels * exp_ratio)
        op_list = kwargs.get('op_list', [])
        self.layers = []

        if norm_layer is not None:
            # Point-wise expansion
            self.conv_pw = nn.Conv2d(in_channels, mid_chs, exp_kernel_size, bias=False)
            self.bn1 = norm_layer(mid_chs)
            self.act1 = nn.ReLU()

            # Depth-wise convolution
            padding = ((stride - 1) + dilation * (dw_kernel_size - 1)) // 2
            self.conv_dw = dwconv(mid_chs, mid_chs, dw_kernel_size, stride=stride, dilation=dilation,
                                        padding=padding, groups=mid_chs, bias=False)
            self.bn2 = norm_layer(mid_chs)
            self.act2 = nn.ReLU()

            # Point-wise linear projection
            self.conv_pwl = nn.Conv2d(mid_chs, out_channels, pw_kernel_size, groups=pw_group_num, bias=False)
            self.bn3 = norm_layer(out_channels)
            self.layers = ['conv_pw', 'bn1', 'act1', 'conv_dw', 'bn2', 'act2', 'conv_pwl', 'bn3']
        else:
            # Point-wise expansion
            self.conv_pw = nn.Conv2d(in_channels, mid_chs, exp_kernel_size, bias=True)
            self.act1 = nn.ReLU()

            # Depth-wise convolution
            padding = ((stride - 1) + dilation * (dw_kernel_size - 1)) // 2
            self.conv_dw = dwconv(mid_chs, mid_chs, dw_kernel_size, stride=stride, dilation=dilation,
                                        padding=padding, groups=mid_chs, bias=True)
            self.act2 = nn.ReLU()

            # Point-wise linear projection
            self.conv_pwl = nn.Conv2d(mid_chs, out_channels, pw_kernel_size, groups=pw_group_num, bias=True)
            self.layers = ['conv_pw', 'act1', 'conv_dw', 'act2', 'conv_pwl']

    def feature_info(self, location):
        if location == 'expansion':  # after SE, input to PWL
            info = dict(module='conv_pwl', hook_type='forward_pre', num_chs=self.in_channels)
        else:  # location == 'bottleneck', block output
            info = dict(module='', hook_type='', num_chs=self.out_channels)
        return info

    def forward(self, x):
        shortcut = x

        for name in self.layers:
            x = getattr(self, name)(x)

        if self.has_residual:
            if self.drop_path_rate > 0.:
                x = drop_path(x, self.drop_path_rate, self.training)
            x += shortcut

        return x
    
@BLOCKS.register
@PARAMS_MAPPING.register
@ABBREVS.register('cn')
class ConvBnAct(nn.Module):
    """ Conv + Norm Layer + Activation w/ optional skip connection
    """
    _PARAMS_DICT = {
        'i': 'in_channels',
        'o': 'out_channels',
        'k': 'kernel_size',
        's': 'stride',
        'skip': 'skip',
        'n': 'act_layer',
        'd': 'dilation'
    }

    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 skip=False, norm_layer='norm', drop_path_rate=0., act_layer='ReLU', dilation=1, **kwargs):
        super(ConvBnAct, self).__init__()
        self.out_chs = out_channels
        self.has_residual = skip and stride == 1 and in_channels == out_channels
        self.drop_path_rate = drop_path_rate
        op_list = kwargs.get('op_list', [])
        self.layers = []

        padding = int(((stride - 1) + dilation * (kernel_size - 1)) // 2)
        if norm_layer is not None:
            self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=False)
            self.bn = nn.BatchNorm2d(out_channels)
            self.layers = ['conv', 'bn', 'act']
        else:
            self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=True)
            self.layers = ['conv', 'act']

        self.act = nn.ReLU()

    def feature_info(self, location):
        if location == 'expansion':  # output of conv after act, same as block coutput
            info = dict(module='act1', hook_type='forward', num_chs=self.out_chs)
        else:  # location == 'bottleneck', block output
            info = dict(module='', hook_type='', num_chs=self.out_chs)
        return info

    def forward(self, x):
        shortcut = x
        for name in self.layers:
            x = getattr(self, name)(x)
        if self.has_residual:
            if self.drop_path_rate > 0.:
                x = drop_path(x, self.drop_path_rate, self.training)
            x += shortcut
        return x