from lib.models.lightnet.inplace_abn.iabn import InPlaceABN
import torch.nn as nn


class InvertedResidual(nn.Module):
    def __init__(self, inp, oup, stride, dilate, expand_ratio):
        """
        InvertedResidual: Core block of the MobileNetV2
        :param inp:    (int) Number of the input channels
        :param oup:    (int) Number of the output channels
        :param stride: (int) Stride used in the Conv3x3
        :param dilate: (int) Dilation used in the Conv3x3
        :param expand_ratio: (int) Expand ratio of the Channel Width of the Block
        """
        super(InvertedResidual, self).__init__()
        self.stride = stride
        assert stride in [1, 2]

        self.use_res_connect = self.stride == 1 and inp == oup

        self.conv = nn.Sequential(
            # step 1. point-wise convolution
            nn.Conv2d(
                in_channels=inp,
                out_channels=inp * expand_ratio,
                kernel_size=1,
                stride=1,
                padding=0,
                dilation=1,
                groups=1,
                bias=False),
            nn.BatchNorm2d(num_features=inp * expand_ratio),
            nn.LeakyReLU(inplace=True, negative_slope=0.01),

            # step 2. depth-wise convolution
            nn.Conv2d(
                in_channels=inp * expand_ratio,
                out_channels=inp * expand_ratio,
                kernel_size=3,
                stride=stride,
                padding=dilate,
                dilation=dilate,
                groups=inp * expand_ratio,
                bias=False),
            nn.BatchNorm2d(num_features=inp * expand_ratio),
            nn.LeakyReLU(inplace=True, negative_slope=0.01),

            # step 3. point-wise convolution
            nn.Conv2d(
                in_channels=inp * expand_ratio,
                out_channels=oup,
                kernel_size=1,
                stride=1,
                padding=0,
                dilation=1,
                groups=1,
                bias=False),
            nn.BatchNorm2d(num_features=oup),
        )

    def forward(self, x):
        if self.use_res_connect:
            return x + self.conv(x)
        else:
            return self.conv(x)


class InvertedResidualIABN(nn.Module):
    def __init__(self,
                 inp,
                 oup,
                 stride,
                 dilate,
                 expand_ratio,
                 norm_act=InPlaceABN):
        """
        InvertedResidual: Core block of the MobileNetV2
        :param inp:    (int) Number of the input channels
        :param oup:    (int) Number of the output channels
        :param stride: (int) Stride used in the Conv3x3
        :param dilate: (int) Dilation used in the Conv3x3
        :param expand_ratio: (int) Expand ratio of the Channel Width of the Block
        """
        super(InvertedResidualIABN, self).__init__()
        self.stride = stride
        assert stride in [1, 2]

        self.use_res_connect = self.stride == 1 and inp == oup

        self.conv = nn.Sequential(
            # step 1. point-wise convolution
            norm_act(inp),
            nn.Conv2d(
                in_channels=inp,
                out_channels=inp * expand_ratio,
                kernel_size=1,
                stride=1,
                padding=0,
                dilation=1,
                groups=1,
                bias=False),

            # step 2. depth-wise convolution
            norm_act(inp * expand_ratio),
            nn.Conv2d(
                in_channels=inp * expand_ratio,
                out_channels=inp * expand_ratio,
                kernel_size=3,
                stride=stride,
                padding=dilate,
                dilation=dilate,
                groups=inp * expand_ratio,
                bias=False),

            # step 3. point-wise convolution
            norm_act(inp * expand_ratio),
            nn.Conv2d(
                in_channels=inp * expand_ratio,
                out_channels=oup,
                kernel_size=1,
                stride=1,
                padding=0,
                dilation=1,
                groups=1,
                bias=False))

    def forward(self, x):
        if self.use_res_connect:
            return x + self.conv(x)
        else:
            return self.conv(x)
