from collections import namedtuple

import torch
from torch.nn import (AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d,
                      Module, PReLU, ReLU, Sequential, Sigmoid)

# yapf: disable
"""
ArcFace implementation from [TreB1eN](https://github.com/TreB1eN/InsightFace_Pytorch) # isort:skip  # noqa
"""
# yapf: enable


class Flatten(Module):
    """Flatten Module."""

    def forward(self, input):
        return input.view(input.size(0), -1)


def l2_norm(input, axis=1):
    """l2 normalization.

    Args:
        input (torch.Tensor): The input tensor.
        axis (int, optional): Specifies which axis of input to calculate the
            norm across. Defaults to 1.

    Returns:
        Tensor: Tensor after L2 normalization per-instance.
    """
    norm = torch.norm(input, 2, axis, True)
    output = torch.div(input, norm)
    return output


class Bottleneck(namedtuple('Block', ['in_channel', 'depth', 'stride'])):
    """A named tuple describing a ResNet block."""


def get_block(in_channel, depth, num_units, stride=2):
    """Get a single block config.

    Args:
        in_channel (int): Input channels.
        depth (int): Output channels.
        num_units (int): Number of unit modules.
        stride (int, optional): Conv2d stride. Defaults to 2.

    Returns:
        list: A list of unit modules' config.
    """
    return [Bottleneck(in_channel, depth, stride)
            ] + [Bottleneck(depth, depth, 1) for i in range(num_units - 1)]


def get_blocks(num_layers):
    """Get block configs of backbone.

    Args:
        num_layers (int): Number of ConvBlock layers in backbone.

    Raises:
        ValueError: `num_layers` must be one of [50, 100, 152].

    Returns:
        list: A list of block configs.
    """
    if num_layers == 50:
        blocks = [
            get_block(in_channel=64, depth=64, num_units=3),
            get_block(in_channel=64, depth=128, num_units=4),
            get_block(in_channel=128, depth=256, num_units=14),
            get_block(in_channel=256, depth=512, num_units=3)
        ]
    elif num_layers == 100:
        blocks = [
            get_block(in_channel=64, depth=64, num_units=3),
            get_block(in_channel=64, depth=128, num_units=13),
            get_block(in_channel=128, depth=256, num_units=30),
            get_block(in_channel=256, depth=512, num_units=3)
        ]
    elif num_layers == 152:
        blocks = [
            get_block(in_channel=64, depth=64, num_units=3),
            get_block(in_channel=64, depth=128, num_units=8),
            get_block(in_channel=128, depth=256, num_units=36),
            get_block(in_channel=256, depth=512, num_units=3)
        ]
    else:
        raise ValueError(
            'Invalid number of layers: {}. Must be one of [50, 100, 152]'.
            format(num_layers))
    return blocks


class SEModule(Module):
    """Squeeze-and-Excitation Modules.

    Args:
        channels (int): Input channels.
        reduction (int): Intermediate channels reduction ratio.
    """

    def __init__(self, channels, reduction):
        super(SEModule, self).__init__()
        self.avg_pool = AdaptiveAvgPool2d(1)
        self.fc1 = Conv2d(
            channels,
            channels // reduction,
            kernel_size=1,
            padding=0,
            bias=False)
        self.relu = ReLU(inplace=True)
        self.fc2 = Conv2d(
            channels // reduction,
            channels,
            kernel_size=1,
            padding=0,
            bias=False)
        self.sigmoid = Sigmoid()

    def forward(self, x):
        """Forward Function."""
        module_input = x
        x = self.avg_pool(x)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.sigmoid(x)
        return module_input * x


class bottleneck_IR(Module):
    """Intermediate Resblock of bottleneck.

    Args:
        in_channel (int): Input channels.
        depth (int): Output channels.
        stride (int): Conv2d stride.
    """

    def __init__(self, in_channel, depth, stride):
        """Intermediate Resblock of bottleneck.

        Args:
            in_channel (int): Input channels.
            depth (int): Output channels.
            stride (int): Conv2d stride.
        """
        super(bottleneck_IR, self).__init__()
        if in_channel == depth:
            self.shortcut_layer = MaxPool2d(1, stride)
        else:
            self.shortcut_layer = Sequential(
                Conv2d(in_channel, depth, (1, 1), stride, bias=False),
                BatchNorm2d(depth))
        self.res_layer = Sequential(
            BatchNorm2d(in_channel),
            Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False),
            PReLU(depth), Conv2d(depth, depth, (3, 3), stride, 1, bias=False),
            BatchNorm2d(depth))

    def forward(self, x):
        """Forward function."""
        shortcut = self.shortcut_layer(x)
        res = self.res_layer(x)
        return res + shortcut


class bottleneck_IR_SE(Module):
    """Intermediate Resblock of bottleneck with SEModule.

    Args:
        in_channel (int): Input channels.
        depth (int): Output channels.
        stride (int): Conv2d stride.
    """

    def __init__(self, in_channel, depth, stride):
        super(bottleneck_IR_SE, self).__init__()
        if in_channel == depth:
            self.shortcut_layer = MaxPool2d(1, stride)
        else:
            self.shortcut_layer = Sequential(
                Conv2d(in_channel, depth, (1, 1), stride, bias=False),
                BatchNorm2d(depth))
        self.res_layer = Sequential(
            BatchNorm2d(in_channel),
            Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False),
            PReLU(depth), Conv2d(depth, depth, (3, 3), stride, 1, bias=False),
            BatchNorm2d(depth), SEModule(depth, 16))

    def forward(self, x):
        """Forward function."""
        shortcut = self.shortcut_layer(x)
        res = self.res_layer(x)
        return res + shortcut

[((890, 944), 'collections.namedtuple', 'namedtuple', (['"""Block"""', "['in_channel', 'depth', 'stride']"], {}), "('Block', ['in_channel', 'depth', 'stride'])\n", (900, 944), False, 'from collections import namedtuple\n'), ((784, 816), 'torch.norm', 'torch.norm', (['input', '(2)', 'axis', '(True)'], {}), '(input, 2, axis, True)\n', (794, 816), False, 'import torch\n'), ((830, 852), 'torch.div', 'torch.div', (['input', 'norm'], {}), '(input, norm)\n', (839, 852), False, 'import torch\n'), ((3144, 3164), 'torch.nn.AdaptiveAvgPool2d', 'AdaptiveAvgPool2d', (['(1)'], {}), '(1)\n', (3161, 3164), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((3184, 3261), 'torch.nn.Conv2d', 'Conv2d', (['channels', '(channels // reduction)'], {'kernel_size': '(1)', 'padding': '(0)', 'bias': '(False)'}), '(channels, channels // reduction, kernel_size=1, padding=0, bias=False)\n', (3190, 3261), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((3343, 3361), 'torch.nn.ReLU', 'ReLU', ([], {'inplace': '(True)'}), '(inplace=True)\n', (3347, 3361), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((3381, 3458), 'torch.nn.Conv2d', 'Conv2d', (['(channels // reduction)', 'channels'], {'kernel_size': '(1)', 'padding': '(0)', 'bias': '(False)'}), '(channels // reduction, channels, kernel_size=1, padding=0, bias=False)\n', (3387, 3458), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((3543, 3552), 'torch.nn.Sigmoid', 'Sigmoid', ([], {}), '()\n', (3550, 3552), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((4378, 4398), 'torch.nn.MaxPool2d', 'MaxPool2d', (['(1)', 'stride'], {}), '(1, stride)\n', (4387, 4398), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((4615, 4638), 'torch.nn.BatchNorm2d', 'BatchNorm2d', (['in_channel'], {}), '(in_channel)\n', (4626, 4638), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((4652, 4708), 'torch.nn.Conv2d', 'Conv2d', (['in_channel', 'depth', '(3, 3)', '(1, 1)', '(1)'], {'bias': '(False)'}), '(in_channel, depth, (3, 3), (1, 1), 1, bias=False)\n', (4658, 4708), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((4722, 4734), 'torch.nn.PReLU', 'PReLU', (['depth'], {}), '(depth)\n', (4727, 4734), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((4736, 4787), 'torch.nn.Conv2d', 'Conv2d', (['depth', 'depth', '(3, 3)', 'stride', '(1)'], {'bias': '(False)'}), '(depth, depth, (3, 3), stride, 1, bias=False)\n', (4742, 4787), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((4801, 4819), 'torch.nn.BatchNorm2d', 'BatchNorm2d', (['depth'], {}), '(depth)\n', (4812, 4819), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((5379, 5399), 'torch.nn.MaxPool2d', 'MaxPool2d', (['(1)', 'stride'], {}), '(1, stride)\n', (5388, 5399), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((5616, 5639), 'torch.nn.BatchNorm2d', 'BatchNorm2d', (['in_channel'], {}), '(in_channel)\n', (5627, 5639), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((5653, 5709), 'torch.nn.Conv2d', 'Conv2d', (['in_channel', 'depth', '(3, 3)', '(1, 1)', '(1)'], {'bias': '(False)'}), '(in_channel, depth, (3, 3), (1, 1), 1, bias=False)\n', (5659, 5709), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((5723, 5735), 'torch.nn.PReLU', 'PReLU', (['depth'], {}), '(depth)\n', (5728, 5735), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((5737, 5788), 'torch.nn.Conv2d', 'Conv2d', (['depth', 'depth', '(3, 3)', 'stride', '(1)'], {'bias': '(False)'}), '(depth, depth, (3, 3), stride, 1, bias=False)\n', (5743, 5788), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((5802, 5820), 'torch.nn.BatchNorm2d', 'BatchNorm2d', (['depth'], {}), '(depth)\n', (5813, 5820), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((4475, 4528), 'torch.nn.Conv2d', 'Conv2d', (['in_channel', 'depth', '(1, 1)', 'stride'], {'bias': '(False)'}), '(in_channel, depth, (1, 1), stride, bias=False)\n', (4481, 4528), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((4546, 4564), 'torch.nn.BatchNorm2d', 'BatchNorm2d', (['depth'], {}), '(depth)\n', (4557, 4564), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((5476, 5529), 'torch.nn.Conv2d', 'Conv2d', (['in_channel', 'depth', '(1, 1)', 'stride'], {'bias': '(False)'}), '(in_channel, depth, (1, 1), stride, bias=False)\n', (5482, 5529), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n'), ((5547, 5565), 'torch.nn.BatchNorm2d', 'BatchNorm2d', (['depth'], {}), '(depth)\n', (5558, 5565), False, 'from torch.nn import AdaptiveAvgPool2d, BatchNorm2d, Conv2d, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid\n')]