"""
UltraLight PyTorch implementation adapted from:
https://github.com/Linzaer/Ultra-Light-Fast-Generic-Face-Detector-1MB

Original implementation by Linzaer
"""

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


class BasicConv(nn.Module):

    def __init__(
        self,
        in_planes,
        out_planes,
        kernel_size,
        stride=1,
        padding=0,
        dilation=1,
        groups=1,
        relu=True,
        bn=True,
    ):
        super(BasicConv, self).__init__()
        self.out_channels = out_planes
        if bn:
            self.conv = nn.Conv2d(
                in_planes,
                out_planes,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                dilation=dilation,
                groups=groups,
                bias=False,
            )
            self.bn = nn.BatchNorm2d(out_planes, eps=1e-5, momentum=0.01, affine=True)
            self.relu = nn.ReLU(inplace=True) if relu else None
        else:
            self.conv = nn.Conv2d(
                in_planes,
                out_planes,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                dilation=dilation,
                groups=groups,
                bias=True,
            )
            self.bn = None
            self.relu = nn.ReLU(inplace=True) if relu else None

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.relu is not None:
            x = self.relu(x)
        return x


class BasicRFB(nn.Module):

    def __init__(
        self,
        in_planes,
        out_planes,
        stride=1,
        scale=0.1,
        map_reduce=8,
        vision=1,
        groups=1,
    ):
        super(BasicRFB, self).__init__()
        self.scale = scale
        self.out_channels = out_planes
        inter_planes = in_planes // map_reduce

        self.branch0 = nn.Sequential(
            BasicConv(
                in_planes,
                inter_planes,
                kernel_size=1,
                stride=1,
                groups=groups,
                relu=False,
            ),
            BasicConv(
                inter_planes,
                2 * inter_planes,
                kernel_size=(3, 3),
                stride=stride,
                padding=(1, 1),
                groups=groups,
            ),
            BasicConv(
                2 * inter_planes,
                2 * inter_planes,
                kernel_size=3,
                stride=1,
                padding=vision + 1,
                dilation=vision + 1,
                relu=False,
                groups=groups,
            ),
        )
        self.branch1 = nn.Sequential(
            BasicConv(
                in_planes,
                inter_planes,
                kernel_size=1,
                stride=1,
                groups=groups,
                relu=False,
            ),
            BasicConv(
                inter_planes,
                2 * inter_planes,
                kernel_size=(3, 3),
                stride=stride,
                padding=(1, 1),
                groups=groups,
            ),
            BasicConv(
                2 * inter_planes,
                2 * inter_planes,
                kernel_size=3,
                stride=1,
                padding=vision + 2,
                dilation=vision + 2,
                relu=False,
                groups=groups,
            ),
        )
        self.branch2 = nn.Sequential(
            BasicConv(
                in_planes,
                inter_planes,
                kernel_size=1,
                stride=1,
                groups=groups,
                relu=False,
            ),
            BasicConv(
                inter_planes,
                (inter_planes // 2) * 3,
                kernel_size=3,
                stride=1,
                padding=1,
                groups=groups,
            ),
            BasicConv(
                (inter_planes // 2) * 3,
                2 * inter_planes,
                kernel_size=3,
                stride=stride,
                padding=1,
                groups=groups,
            ),
            BasicConv(
                2 * inter_planes,
                2 * inter_planes,
                kernel_size=3,
                stride=1,
                padding=vision + 4,
                dilation=vision + 4,
                relu=False,
                groups=groups,
            ),
        )

        self.ConvLinear = BasicConv(
            6 * inter_planes, out_planes, kernel_size=1, stride=1, relu=False
        )
        self.shortcut = BasicConv(
            in_planes, out_planes, kernel_size=1, stride=stride, relu=False
        )
        self.relu = nn.ReLU(inplace=False)

    def forward(self, x):
        x0 = self.branch0(x)
        x1 = self.branch1(x)
        x2 = self.branch2(x)

        out = torch.cat((x0, x1, x2), 1)
        out = self.ConvLinear(out)
        short = self.shortcut(x)
        out = out * self.scale + short
        out = self.relu(out)

        return out


class Mb_Tiny_RFB(nn.Module):

    def __init__(self, num_classes=2):
        super(Mb_Tiny_RFB, self).__init__()
        self.base_channel = 8 * 2

        def conv_bn(inp, oup, stride):
            return nn.Sequential(
                nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
                nn.BatchNorm2d(oup),
                nn.ReLU(inplace=True),
            )

        def conv_dw(inp, oup, stride):
            return nn.Sequential(
                nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
                nn.BatchNorm2d(inp),
                nn.ReLU(inplace=True),
                nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
                nn.ReLU(inplace=True),
            )

        self.model = nn.Sequential(
            conv_bn(3, self.base_channel, 2),  # 160*120
            conv_dw(self.base_channel, self.base_channel * 2, 1),
            conv_dw(self.base_channel * 2, self.base_channel * 2, 2),  # 80*60
            conv_dw(self.base_channel * 2, self.base_channel * 2, 1),
            conv_dw(self.base_channel * 2, self.base_channel * 4, 2),  # 40*30
            conv_dw(self.base_channel * 4, self.base_channel * 4, 1),
            conv_dw(self.base_channel * 4, self.base_channel * 4, 1),
            BasicRFB(self.base_channel * 4, self.base_channel * 4, stride=1, scale=1.0),
            conv_dw(self.base_channel * 4, self.base_channel * 8, 2),  # 20*15
            conv_dw(self.base_channel * 8, self.base_channel * 8, 1),
            conv_dw(self.base_channel * 8, self.base_channel * 8, 1),
            conv_dw(self.base_channel * 8, self.base_channel * 16, 2),  # 10*8
            conv_dw(self.base_channel * 16, self.base_channel * 16, 1),
        )
        self.fc = nn.Linear(1024, num_classes)

    def forward(self, x):
        x = self.model(x)
        x = F.avg_pool2d(x, 7)
        x = x.view(-1, 1024)
        x = self.fc(x)
        return x
