# Copyright 2021 The KaiJIN Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.init as init
from .box_utils import Detect, PriorBox
import time


class FEM(nn.Module):

  def __init__(self, in_planes):
    super(FEM, self).__init__()
    inter_planes = in_planes // 3
    inter_planes1 = in_planes - 2 * inter_planes
    self.branch1 = nn.Conv2d(in_planes, inter_planes, kernel_size=3, stride=1, padding=3, dilation=3)

    self.branch2 = nn.Sequential(
        nn.Conv2d(in_planes, inter_planes, kernel_size=3, stride=1, padding=3, dilation=3),
        nn.ReLU(inplace=True),
        nn.Conv2d(inter_planes, inter_planes, kernel_size=3, stride=1, padding=3, dilation=3)
    )
    self.branch3 = nn.Sequential(
        nn.Conv2d(in_planes, inter_planes1, kernel_size=3, stride=1, padding=3, dilation=3),
        nn.ReLU(inplace=True),
        nn.Conv2d(inter_planes1, inter_planes1, kernel_size=3, stride=1, padding=3, dilation=3),
        nn.ReLU(inplace=True),
        nn.Conv2d(inter_planes1, inter_planes1, kernel_size=3, stride=1, padding=3, dilation=3)
    )

  def forward(self, x):
    x1 = self.branch1(x)
    x2 = self.branch2(x)
    x3 = self.branch3(x)
    out = torch.cat((x1, x2, x3), dim=1)
    out = F.relu(out, inplace=True)
    return out


class L2Norm(nn.Module):

  def __init__(self, n_channels, scale):
    super(L2Norm, self).__init__()
    self.n_channels = n_channels
    self.gamma = scale or None
    self.eps = 1e-10
    self.weight = nn.Parameter(torch.Tensor(self.n_channels))
    self.reset_parameters()

  def reset_parameters(self):
    init.constant_(self.weight, self.gamma)

  def forward(self, x):
    norm = x.pow(2).sum(dim=1, keepdim=True).sqrt() + self.eps
    x = torch.div(x, norm)
    out = self.weight.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand_as(x) * x
    return out


class DSFDNet(nn.Module):

  def __init__(self, device='cuda'):
    super(DSFDNet, self).__init__()
    self.device = device

    self.vgg = nn.ModuleList([
        nn.Conv2d(3, 64, 3, 1, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(64, 64, 3, 1, padding=1),
        nn.ReLU(inplace=True),
        nn.MaxPool2d(2, 2),

        nn.Conv2d(64, 128, 3, 1, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(128, 128, 3, 1, padding=1),
        nn.ReLU(inplace=True),
        nn.MaxPool2d(2, 2),

        nn.Conv2d(128, 256, 3, 1, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(256, 256, 3, 1, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(256, 256, 3, 1, padding=1),
        nn.ReLU(inplace=True),
        nn.MaxPool2d(2, 2, ceil_mode=True),

        nn.Conv2d(256, 512, 3, 1, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(512, 512, 3, 1, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(512, 512, 3, 1, padding=1),
        nn.ReLU(inplace=True),
        nn.MaxPool2d(2, 2),

        nn.Conv2d(512, 512, 3, 1, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(512, 512, 3, 1, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(512, 512, 3, 1, padding=1),
        nn.ReLU(inplace=True),
        nn.MaxPool2d(2, 2),

        nn.Conv2d(512, 1024, 3, 1, padding=6, dilation=6),
        nn.ReLU(inplace=True),
        nn.Conv2d(1024, 1024, 1, 1),
        nn.ReLU(inplace=True),
    ])

    self.L2Normof1 = L2Norm(256, 10)
    self.L2Normof2 = L2Norm(512, 8)
    self.L2Normof3 = L2Norm(512, 5)

    self.extras = nn.ModuleList([
        nn.Conv2d(1024, 256, 1, 1),
        nn.Conv2d(256, 512, 3, 2, padding=1),
        nn.Conv2d(512, 128, 1, 1),
        nn.Conv2d(128, 256, 3, 2, padding=1),
    ])

    self.fpn_topdown = nn.ModuleList([
        nn.Conv2d(256, 256, 1, 1, padding=0),
        nn.Conv2d(256, 512, 1, 1, padding=0),
        nn.Conv2d(512, 1024, 1, 1, padding=0),
        nn.Conv2d(1024, 512, 1, 1, padding=0),
        nn.Conv2d(512, 512, 1, 1, padding=0),
        nn.Conv2d(512, 256, 1, 1, padding=0),
    ])

    self.fpn_latlayer = nn.ModuleList([
        nn.Conv2d(512, 512, 1, 1, padding=0),
        nn.Conv2d(1024, 1024, 1, 1, padding=0),
        nn.Conv2d(512, 512, 1, 1, padding=0),
        nn.Conv2d(512, 512, 1, 1, padding=0),
        nn.Conv2d(256, 256, 1, 1, padding=0),
    ])

    self.fpn_fem = nn.ModuleList([
        FEM(256), FEM(512), FEM(512),
        FEM(1024), FEM(512), FEM(256),
    ])

    self.L2Normef1 = L2Norm(256, 10)
    self.L2Normef2 = L2Norm(512, 8)
    self.L2Normef3 = L2Norm(512, 5)

    self.loc_pal1 = nn.ModuleList([
        nn.Conv2d(256, 4, 3, 1, padding=1),
        nn.Conv2d(512, 4, 3, 1, padding=1),
        nn.Conv2d(512, 4, 3, 1, padding=1),
        nn.Conv2d(1024, 4, 3, 1, padding=1),
        nn.Conv2d(512, 4, 3, 1, padding=1),
        nn.Conv2d(256, 4, 3, 1, padding=1),
    ])

    self.conf_pal1 = nn.ModuleList([
        nn.Conv2d(256, 2, 3, 1, padding=1),
        nn.Conv2d(512, 2, 3, 1, padding=1),
        nn.Conv2d(512, 2, 3, 1, padding=1),
        nn.Conv2d(1024, 2, 3, 1, padding=1),
        nn.Conv2d(512, 2, 3, 1, padding=1),
        nn.Conv2d(256, 2, 3, 1, padding=1),
    ])

    self.loc_pal2 = nn.ModuleList([
        nn.Conv2d(256, 4, 3, 1, padding=1),
        nn.Conv2d(512, 4, 3, 1, padding=1),
        nn.Conv2d(512, 4, 3, 1, padding=1),
        nn.Conv2d(1024, 4, 3, 1, padding=1),
        nn.Conv2d(512, 4, 3, 1, padding=1),
        nn.Conv2d(256, 4, 3, 1, padding=1),
    ])

    self.conf_pal2 = nn.ModuleList([
        nn.Conv2d(256, 2, 3, 1, padding=1),
        nn.Conv2d(512, 2, 3, 1, padding=1),
        nn.Conv2d(512, 2, 3, 1, padding=1),
        nn.Conv2d(1024, 2, 3, 1, padding=1),
        nn.Conv2d(512, 2, 3, 1, padding=1),
        nn.Conv2d(256, 2, 3, 1, padding=1),
    ])

    self.softmax = nn.Softmax(dim=-1)
    self.detect = Detect()

  def _upsample_prod(self, x, y):
    _, _, H, W = y.size()
    # return F.upsample(x, size=(H, W), mode='bilinear') * y
    return F.interpolate(x, size=(H, W), mode='bilinear', align_corners=True) * y

  def forward(self, x):
    size = x.size()[2:]
    pal1_sources = list()
    pal2_sources = list()
    loc_pal1 = list()
    conf_pal1 = list()
    loc_pal2 = list()
    conf_pal2 = list()

    # vgg layers
    for k in range(16):
      x = self.vgg[k](x)
    of1 = x
    s = self.L2Normof1(of1)
    pal1_sources.append(s)

    for k in range(16, 23):
      x = self.vgg[k](x)
    of2 = x
    s = self.L2Normof2(of2)
    pal1_sources.append(s)

    for k in range(23, 30):
      x = self.vgg[k](x)
    of3 = x
    s = self.L2Normof3(of3)
    pal1_sources.append(s)

    for k in range(30, len(self.vgg)):
      x = self.vgg[k](x)
    of4 = x
    pal1_sources.append(of4)

    # extra layers
    for k in range(2):
      x = F.relu(self.extras[k](x), inplace=True)
    of5 = x
    pal1_sources.append(of5)

    for k in range(2, 4):
      x = F.relu(self.extras[k](x), inplace=True)
    of6 = x
    pal1_sources.append(of6)

    conv7 = F.relu(self.fpn_topdown[0](of6), inplace=True)

    x = F.relu(self.fpn_topdown[1](conv7), inplace=True)
    conv6 = F.relu(self._upsample_prod(x, self.fpn_latlayer[0](of5)), inplace=True)

    x = F.relu(self.fpn_topdown[2](conv6), inplace=True)
    convfc7_2 = F.relu(self._upsample_prod(x, self.fpn_latlayer[1](of4)), inplace=True)

    x = F.relu(self.fpn_topdown[3](convfc7_2), inplace=True)
    conv5 = F.relu(self._upsample_prod(x, self.fpn_latlayer[2](of3)), inplace=True)

    x = F.relu(self.fpn_topdown[4](conv5), inplace=True)
    conv4 = F.relu(self._upsample_prod(x, self.fpn_latlayer[3](of2)), inplace=True)

    x = F.relu(self.fpn_topdown[5](conv4), inplace=True)
    conv3 = F.relu(self._upsample_prod(x, self.fpn_latlayer[4](of1)), inplace=True)

    ef1 = self.fpn_fem[0](conv3)
    ef1 = self.L2Normef1(ef1)
    ef2 = self.fpn_fem[1](conv4)
    ef2 = self.L2Normef2(ef2)
    ef3 = self.fpn_fem[2](conv5)
    ef3 = self.L2Normef3(ef3)
    ef4 = self.fpn_fem[3](convfc7_2)
    ef5 = self.fpn_fem[4](conv6)
    ef6 = self.fpn_fem[5](conv7)

    pal2_sources = (ef1, ef2, ef3, ef4, ef5, ef6)

    for (x, l, c) in zip(pal1_sources, self.loc_pal1, self.conf_pal1):
      loc_pal1.append(l(x).permute(0, 2, 3, 1).contiguous())
      conf_pal1.append(c(x).permute(0, 2, 3, 1).contiguous())

    for (x, l, c) in zip(pal2_sources, self.loc_pal2, self.conf_pal2):
      loc_pal2.append(l(x).permute(0, 2, 3, 1).contiguous())
      conf_pal2.append(c(x).permute(0, 2, 3, 1).contiguous())

    features_maps = []
    for i in range(len(loc_pal1)):
      feat = []
      feat += [loc_pal1[i].size(1), loc_pal1[i].size(2)]
      features_maps += [feat]

    loc_pal1 = torch.cat([o.view(o.size(0), -1) for o in loc_pal1], 1)
    conf_pal1 = torch.cat([o.view(o.size(0), -1) for o in conf_pal1], 1)

    loc_pal2 = torch.cat([o.view(o.size(0), -1) for o in loc_pal2], 1)
    conf_pal2 = torch.cat([o.view(o.size(0), -1) for o in conf_pal2], 1)

    with torch.no_grad():
      self.priorbox1 = PriorBox(size, features_maps, pal=1)
      self.priors_pal1 = self.priorbox1.forward()
      self.priorbox2 = PriorBox(size, features_maps, pal=2)
      self.priors_pal2 = self.priorbox2.forward()

    output = self.detect(
        loc_pal2.view(loc_pal2.size(0), -1, 4),
        self.softmax(conf_pal2.view(conf_pal2.size(0), -1, 2)),
        self.priors_pal2.type(type(x.data)).to(self.device)
    )

    return output
