"""
@author: Yibo Hu, Jun Wang
@date: 20201019
@contact: jun21wangustc@gmail.com
"""

from collections import OrderedDict

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


def channel_shuffle(x, groups):
    assert groups > 1
    batchsize, num_channels, height, width = x.size()
    assert num_channels % groups == 0
    channels_per_group = num_channels // groups
    # reshape
    x = x.view(batchsize, groups, channels_per_group, height, width)
    # transpose
    x = torch.transpose(x, 1, 2).contiguous()
    # flatten
    x = x.view(batchsize, -1, height, width)
    return x


def get_same_padding(kernel_size):
    if isinstance(kernel_size, tuple):
        assert len(kernel_size) == 2, "invalid kernel size: {}".format(
            kernel_size
        )
        p1 = get_same_padding(kernel_size[0])
        p2 = get_same_padding(kernel_size[1])
        return p1, p2
    assert isinstance(
        kernel_size, int
    ), "kernel size should be either `int` or `tuple`"
    assert kernel_size % 2 > 0, "kernel size should be odd number"
    return kernel_size // 2


class Swish(nn.Module):
    def __init__(self, inplace=False):
        super(Swish, self).__init__()
        self.inplace = inplace

    def forward(self, x):
        if self.inplace:
            return x.mul_(x.sigmoid())
        else:
            return x * x.sigmoid()


class HardSwish(nn.Module):
    def __init__(self, inplace=False):
        super(HardSwish, self).__init__()
        self.inplace = inplace

    def forward(self, x):
        if self.inplace:
            return x.mul_(F.relu6(x + 3.0, inplace=True) / 6.0)
        else:
            return x * F.relu6(x + 3.0) / 6.0


class BasicLayer(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        use_bn=True,
        affine=True,
        act_func="relu6",
        ops_order="weight_bn_act",
    ):
        super(BasicLayer, self).__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.use_bn = use_bn
        self.affine = affine
        self.act_func = act_func
        self.ops_order = ops_order

        """ add modules """
        # batch norm
        if self.use_bn:
            if self.bn_before_weight:
                self.bn = nn.BatchNorm2d(
                    in_channels, affine=affine, track_running_stats=affine
                )
            else:
                self.bn = nn.BatchNorm2d(
                    out_channels, affine=affine, track_running_stats=affine
                )
        else:
            self.bn = None
        # activation
        if act_func == "relu":
            if self.ops_list[0] == "act":
                self.act = nn.ReLU(inplace=False)
            else:
                self.act = nn.ReLU(inplace=True)
        elif act_func == "relu6":
            if self.ops_list[0] == "act":
                self.act = nn.ReLU6(inplace=False)
            else:
                self.act = nn.ReLU6(inplace=True)
        elif act_func == "swish":
            if self.ops_list[0] == "act":
                self.act = Swish(inplace=False)
            else:
                self.act = Swish(inplace=True)
        elif act_func == "h-swish":
            if self.ops_list[0] == "act":
                self.act = HardSwish(inplace=False)
            else:
                self.act = HardSwish(inplace=True)
        else:
            self.act = None

    @property
    def ops_list(self):
        return self.ops_order.split("_")

    @property
    def bn_before_weight(self):
        for op in self.ops_list:
            if op == "bn":
                return True
            elif op == "weight":
                return False
        raise ValueError("Invalid ops_order: %s" % self.ops_order)

    def weight_call(self, x):
        raise NotImplementedError

    def forward(self, x):
        for op in self.ops_list:
            if op == "weight":
                x = self.weight_call(x)
            elif op == "bn":
                if self.bn is not None:
                    x = self.bn(x)
            elif op == "act":
                if self.act is not None:
                    x = self.act(x)
            else:
                raise ValueError("Unrecognized op: %s" % op)
        return x


class ConvLayer(BasicLayer):
    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size=3,
        stride=1,
        groups=1,
        has_shuffle=False,
        bias=False,
        use_bn=True,
        affine=True,
        act_func="relu6",
        ops_order="weight_bn_act",
    ):
        super(ConvLayer, self).__init__(
            in_channels, out_channels, use_bn, affine, act_func, ops_order
        )

        self.kernel_size = kernel_size
        self.stride = stride
        self.groups = groups
        self.has_shuffle = has_shuffle
        self.bias = bias

        padding = get_same_padding(self.kernel_size)
        self.conv = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size=self.kernel_size,
            stride=self.stride,
            padding=padding,
            groups=self.groups,
            bias=self.bias,
        )

    def weight_call(self, x):
        x = self.conv(x)
        if self.has_shuffle and self.groups > 1:
            x = channel_shuffle(x, self.groups)
        return x


class LinearLayer(nn.Module):
    def __init__(
        self,
        in_features,
        out_features,
        bias=True,
        use_bn=False,
        affine=False,
        act_func=None,
        ops_order="weight_bn_act",
    ):
        super(LinearLayer, self).__init__()

        self.in_features = in_features
        self.out_features = out_features
        self.bias = bias
        self.use_bn = use_bn
        self.affine = affine
        self.act_func = act_func
        self.ops_order = ops_order

        """ add modules """
        # batch norm
        if self.use_bn:
            if self.bn_before_weight:
                self.bn = nn.BatchNorm1d(
                    in_features, affine=affine, track_running_stats=affine
                )
            else:
                self.bn = nn.BatchNorm1d(
                    out_features, affine=affine, track_running_stats=affine
                )
        else:
            self.bn = None
        # activation
        if act_func == "relu":
            if self.ops_list[0] == "act":
                self.act = nn.ReLU(inplace=False)
            else:
                self.act = nn.ReLU(inplace=True)
        elif act_func == "relu6":
            if self.ops_list[0] == "act":
                self.act = nn.ReLU6(inplace=False)
            else:
                self.act = nn.ReLU6(inplace=True)
        elif act_func == "tanh":
            self.act = nn.Tanh()
        elif act_func == "sigmoid":
            self.act = nn.Sigmoid()
        else:
            self.act = None
        # linear
        self.linear = nn.Linear(self.in_features, self.out_features, self.bias)

    @property
    def ops_list(self):
        return self.ops_order.split("_")

    @property
    def bn_before_weight(self):
        for op in self.ops_list:
            if op == "bn":
                return True
            elif op == "weight":
                return False
        raise ValueError("Invalid ops_order: %s" % self.ops_order)

    def forward(self, x):
        for op in self.ops_list:
            if op == "weight":
                x = self.linear(x)
            elif op == "bn":
                if self.bn is not None:
                    x = self.bn(x)
            elif op == "act":
                if self.act is not None:
                    x = self.act(x)
            else:
                raise ValueError("Unrecognized op: %s" % op)
        return x


class MBInvertedResBlock(nn.Module):
    def __init__(
        self,
        in_channels,
        mid_channels,
        se_channels,
        out_channels,
        kernel_size=3,
        stride=1,
        groups=1,
        has_shuffle=False,
        bias=False,
        use_bn=True,
        affine=True,
        act_func="relu6",
    ):
        super(MBInvertedResBlock, self).__init__()

        self.in_channels = in_channels
        self.mid_channels = mid_channels
        self.se_channels = se_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.groups = groups
        self.has_shuffle = has_shuffle
        self.bias = bias
        self.use_bn = use_bn
        self.affine = affine
        self.act_func = act_func

        # inverted bottleneck
        if mid_channels > in_channels:
            inverted_bottleneck = OrderedDict(
                [
                    (
                        "conv",
                        nn.Conv2d(
                            in_channels,
                            mid_channels,
                            1,
                            1,
                            0,
                            groups=groups,
                            bias=bias,
                        ),
                    ),
                ]
            )
            if use_bn:
                inverted_bottleneck["bn"] = nn.BatchNorm2d(
                    mid_channels, affine=affine, track_running_stats=affine
                )
            if act_func == "relu":
                inverted_bottleneck["act"] = nn.ReLU(inplace=True)
            elif act_func == "relu6":
                inverted_bottleneck["act"] = nn.ReLU6(inplace=True)
            elif act_func == "swish":
                inverted_bottleneck["act"] = Swish(inplace=True)
            elif act_func == "h-swish":
                inverted_bottleneck["act"] = HardSwish(inplace=True)
            self.inverted_bottleneck = nn.Sequential(inverted_bottleneck)
        else:
            self.inverted_bottleneck = None
            self.mid_channels = in_channels
            mid_channels = in_channels

        # depthwise convolution
        padding = get_same_padding(self.kernel_size)
        depth_conv = OrderedDict(
            [
                (
                    "conv",
                    nn.Conv2d(
                        mid_channels,
                        mid_channels,
                        kernel_size,
                        stride,
                        padding,
                        groups=mid_channels,
                        bias=bias,
                    ),
                ),
            ]
        )
        if use_bn:
            depth_conv["bn"] = nn.BatchNorm2d(
                mid_channels, affine=affine, track_running_stats=affine
            )
        if act_func == "relu":
            depth_conv["act"] = nn.ReLU(inplace=True)
        elif act_func == "relu6":
            depth_conv["act"] = nn.ReLU6(inplace=True)
        elif act_func == "swish":
            depth_conv["act"] = Swish(inplace=True)
        elif act_func == "h-swish":
            depth_conv["act"] = HardSwish(inplace=True)
        self.depth_conv = nn.Sequential(depth_conv)

        # se model
        if se_channels > 0:
            squeeze_excite = OrderedDict(
                [
                    (
                        "conv_reduce",
                        nn.Conv2d(
                            mid_channels,
                            se_channels,
                            1,
                            1,
                            0,
                            groups=groups,
                            bias=True,
                        ),
                    ),
                ]
            )
            if act_func == "relu":
                squeeze_excite["act"] = nn.ReLU(inplace=True)
            elif act_func == "relu6":
                squeeze_excite["act"] = nn.ReLU6(inplace=True)
            elif act_func == "swish":
                squeeze_excite["act"] = Swish(inplace=True)
            elif act_func == "h-swish":
                squeeze_excite["act"] = HardSwish(inplace=True)
            squeeze_excite["conv_expand"] = nn.Conv2d(
                se_channels, mid_channels, 1, 1, 0, groups=groups, bias=True
            )
            self.squeeze_excite = nn.Sequential(squeeze_excite)
        else:
            self.squeeze_excite = None
            self.se_channels = 0

        # pointwise linear
        point_linear = OrderedDict(
            [
                (
                    "conv",
                    nn.Conv2d(
                        mid_channels,
                        out_channels,
                        1,
                        1,
                        0,
                        groups=groups,
                        bias=bias,
                    ),
                ),
            ]
        )
        if use_bn:
            point_linear["bn"] = nn.BatchNorm2d(
                out_channels, affine=affine, track_running_stats=affine
            )
        self.point_linear = nn.Sequential(point_linear)

        # residual flag
        self.has_residual = (in_channels == out_channels) and (stride == 1)

    def forward(self, x):
        res = x

        if self.inverted_bottleneck is not None:
            x = self.inverted_bottleneck(x)
            if self.has_shuffle and self.groups > 1:
                x = channel_shuffle(x, self.groups)

        x = self.depth_conv(x)
        if self.squeeze_excite is not None:
            x_se = F.adaptive_avg_pool2d(x, 1)
            x = x * torch.sigmoid(self.squeeze_excite(x_se))

        x = self.point_linear(x)
        if self.has_shuffle and self.groups > 1:
            x = channel_shuffle(x, self.groups)

        if self.has_residual:
            x += res

        return x


class Flatten(nn.Module):
    def forward(self, x):
        return x.view(x.size(0), -1)


class TF_NAS_A(nn.Module):
    def __init__(self, out_h, out_w, feat_dim, drop_ratio=0.0):
        super(TF_NAS_A, self).__init__()
        self.drop_ratio = drop_ratio

        self.first_stem = ConvLayer(
            3, 32, kernel_size=3, stride=1, act_func="relu"
        )
        self.second_stem = MBInvertedResBlock(
            32, 32, 8, 16, kernel_size=3, stride=1, act_func="relu"
        )
        self.stage1 = nn.Sequential(
            MBInvertedResBlock(
                16, 83, 32, 24, kernel_size=3, stride=2, act_func="relu"
            ),
            MBInvertedResBlock(
                24, 128, 0, 24, kernel_size=5, stride=1, act_func="relu"
            ),
        )
        self.stage2 = nn.Sequential(
            MBInvertedResBlock(
                24, 138, 48, 40, kernel_size=3, stride=2, act_func="swish"
            ),
            MBInvertedResBlock(
                40, 297, 0, 40, kernel_size=3, stride=1, act_func="swish"
            ),
            MBInvertedResBlock(
                40, 170, 80, 40, kernel_size=5, stride=1, act_func="swish"
            ),
        )
        self.stage3 = nn.Sequential(
            MBInvertedResBlock(
                40, 248, 80, 80, kernel_size=5, stride=2, act_func="swish"
            ),
            MBInvertedResBlock(
                80, 500, 0, 80, kernel_size=3, stride=1, act_func="swish"
            ),
            MBInvertedResBlock(
                80, 424, 0, 80, kernel_size=3, stride=1, act_func="swish"
            ),
            MBInvertedResBlock(
                80, 477, 0, 80, kernel_size=3, stride=1, act_func="swish"
            ),
        )
        self.stage4 = nn.Sequential(
            MBInvertedResBlock(
                80, 504, 160, 112, kernel_size=3, stride=1, act_func="swish"
            ),
            MBInvertedResBlock(
                112, 796, 0, 112, kernel_size=3, stride=1, act_func="swish"
            ),
            MBInvertedResBlock(
                112, 723, 224, 112, kernel_size=3, stride=1, act_func="swish"
            ),
            MBInvertedResBlock(
                112, 555, 224, 112, kernel_size=3, stride=1, act_func="swish"
            ),
        )
        self.stage5 = nn.Sequential(
            MBInvertedResBlock(
                112, 813, 0, 192, kernel_size=3, stride=2, act_func="swish"
            ),
            MBInvertedResBlock(
                192, 1370, 0, 192, kernel_size=3, stride=1, act_func="swish"
            ),
            MBInvertedResBlock(
                192, 1138, 384, 192, kernel_size=3, stride=1, act_func="swish"
            ),
            MBInvertedResBlock(
                192, 1359, 384, 192, kernel_size=3, stride=1, act_func="swish"
            ),
        )
        self.stage6 = nn.Sequential(
            MBInvertedResBlock(
                192, 1203, 384, 320, kernel_size=5, stride=1, act_func="swish"
            ),
        )
        self.feature_mix_layer = ConvLayer(
            320, 1280, kernel_size=1, stride=1, act_func="none"
        )
        self.output_layer = nn.Sequential(
            nn.Dropout(self.drop_ratio),
            Flatten(),
            nn.Linear(1280 * out_h * out_w, feat_dim),
            nn.BatchNorm1d(feat_dim),
        )

        self._initialization()

    def forward(self, x):
        x = self.first_stem(x)
        x = self.second_stem(x)
        for block in self.stage1:
            x = block(x)
        for block in self.stage2:
            x = block(x)
        for block in self.stage3:
            x = block(x)
        for block in self.stage4:
            x = block(x)
        for block in self.stage5:
            x = block(x)
        for block in self.stage6:
            x = block(x)
        x = self.feature_mix_layer(x)
        x = self.output_layer(x)
        return x

    def _initialization(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                if m.weight is not None:
                    nn.init.constant_(m.weight, 1)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)


if __name__ == "__main__":
    x = torch.rand((2, 3, 112, 112))
    net = TF_NAS_A(7, 7, 512, drop_ratio=0.0)

    x = x.cuda()
    net = net.cuda()

    out = net(x)
    print(out.size())
