# -*- coding: utf-8 -*-
# @Time    : 2018-9-13 16:03
# @Author  : xylon

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

from utils.image_utils import soft_nms_3d, soft_max_and_argmax_1d
from utils.math_utils import L2Norm
from .rf_det_module import RFDetModule


class RFDetSO(RFDetModule):
    def __init__(
        self,
        score_com_strength,
        scale_com_strength,
        nms_thresh,
        nms_ksize,
        topk,
        gauss_ksize,
        gauss_sigma,
        ksize,
        padding,
        dilation,
        scale_list,
    ):
        super(RFDetSO, self).__init__(
            score_com_strength,
            scale_com_strength,
            nms_thresh,
            nms_ksize,
            topk,
            gauss_ksize,
            gauss_sigma,
            ksize,
            padding,
            dilation,
            scale_list,
        )

        self.conv_o3 = nn.Conv2d(
            in_channels=16, out_channels=2, kernel_size=1, stride=1, padding=0
        )
        self.conv_o5 = nn.Conv2d(
            in_channels=16, out_channels=2, kernel_size=1, stride=1, padding=0
        )
        self.conv_o7 = nn.Conv2d(
            in_channels=16, out_channels=2, kernel_size=1, stride=1, padding=0
        )
        self.conv_o9 = nn.Conv2d(
            in_channels=16, out_channels=2, kernel_size=1, stride=1, padding=0
        )
        self.conv_o11 = nn.Conv2d(
            in_channels=16, out_channels=2, kernel_size=1, stride=1, padding=0
        )
        self.conv_o13 = nn.Conv2d(
            in_channels=16, out_channels=2, kernel_size=1, stride=1, padding=0
        )
        self.conv_o15 = nn.Conv2d(
            in_channels=16, out_channels=2, kernel_size=1, stride=1, padding=0
        )
        self.conv_o17 = nn.Conv2d(
            in_channels=16, out_channels=2, kernel_size=1, stride=1, padding=0
        )
        self.conv_o19 = nn.Conv2d(
            in_channels=16, out_channels=2, kernel_size=1, stride=1, padding=0
        )
        self.conv_o21 = nn.Conv2d(
            in_channels=16, out_channels=2, kernel_size=1, stride=1, padding=0
        )

    def forward(self, photos):

        # Extract score map in scale space from 3 to 21
        score_featmaps_s3 = F.leaky_relu(self.insnorm1(self.conv1(photos)))
        score_map_s3 = self.insnorm_s3(self.conv_s3(score_featmaps_s3)).permute(
            0, 2, 3, 1
        )
        orint_map_s3 = (
            L2Norm(self.conv_o3(score_featmaps_s3), dim=1)
            .permute(0, 2, 3, 1)
            .unsqueeze(-2)
        )

        score_featmaps_s5 = F.leaky_relu(self.insnorm2(self.conv2(score_featmaps_s3)))
        score_map_s5 = self.insnorm_s5(self.conv_s5(score_featmaps_s5)).permute(
            0, 2, 3, 1
        )
        orint_map_s5 = (
            L2Norm(self.conv_o5(score_featmaps_s5), dim=1)
            .permute(0, 2, 3, 1)
            .unsqueeze(-2)
        )
        score_featmaps_s5 = score_featmaps_s5 + score_featmaps_s3

        score_featmaps_s7 = F.leaky_relu(self.insnorm3(self.conv3(score_featmaps_s5)))
        score_map_s7 = self.insnorm_s7(self.conv_s7(score_featmaps_s7)).permute(
            0, 2, 3, 1
        )
        orint_map_s7 = (
            L2Norm(self.conv_o7(score_featmaps_s7), dim=1)
            .permute(0, 2, 3, 1)
            .unsqueeze(-2)
        )
        score_featmaps_s7 = score_featmaps_s7 + score_featmaps_s5

        score_featmaps_s9 = F.leaky_relu(self.insnorm4(self.conv4(score_featmaps_s7)))
        score_map_s9 = self.insnorm_s9(self.conv_s9(score_featmaps_s9)).permute(
            0, 2, 3, 1
        )
        orint_map_s9 = (
            L2Norm(self.conv_o9(score_featmaps_s9), dim=1)
            .permute(0, 2, 3, 1)
            .unsqueeze(-2)
        )
        score_featmaps_s9 = score_featmaps_s9 + score_featmaps_s7

        score_featmaps_s11 = F.leaky_relu(self.insnorm5(self.conv5(score_featmaps_s9)))
        score_map_s11 = self.insnorm_s11(self.conv_s11(score_featmaps_s11)).permute(
            0, 2, 3, 1
        )
        orint_map_s11 = (
            L2Norm(self.conv_o11(score_featmaps_s11), dim=1)
            .permute(0, 2, 3, 1)
            .unsqueeze(-2)
        )
        score_featmaps_s11 = score_featmaps_s11 + score_featmaps_s9

        score_featmaps_s13 = F.leaky_relu(self.insnorm6(self.conv6(score_featmaps_s11)))
        score_map_s13 = self.insnorm_s13(self.conv_s13(score_featmaps_s13)).permute(
            0, 2, 3, 1
        )
        orint_map_s13 = (
            L2Norm(self.conv_o13(score_featmaps_s13), dim=1)
            .permute(0, 2, 3, 1)
            .unsqueeze(-2)
        )
        score_featmaps_s13 = score_featmaps_s13 + score_featmaps_s11

        score_featmaps_s15 = F.leaky_relu(self.insnorm7(self.conv7(score_featmaps_s13)))
        score_map_s15 = self.insnorm_s15(self.conv_s15(score_featmaps_s15)).permute(
            0, 2, 3, 1
        )
        orint_map_s15 = (
            L2Norm(self.conv_o15(score_featmaps_s15), dim=1)
            .permute(0, 2, 3, 1)
            .unsqueeze(-2)
        )
        score_featmaps_s15 = score_featmaps_s15 + score_featmaps_s13

        score_featmaps_s17 = F.leaky_relu(self.insnorm8(self.conv8(score_featmaps_s15)))
        score_map_s17 = self.insnorm_s17(self.conv_s17(score_featmaps_s17)).permute(
            0, 2, 3, 1
        )
        orint_map_s17 = (
            L2Norm(self.conv_o17(score_featmaps_s17), dim=1)
            .permute(0, 2, 3, 1)
            .unsqueeze(-2)
        )
        score_featmaps_s17 = score_featmaps_s17 + score_featmaps_s15

        score_featmaps_s19 = F.leaky_relu(self.insnorm9(self.conv9(score_featmaps_s17)))
        score_map_s19 = self.insnorm_s19(self.conv_s19(score_featmaps_s19)).permute(
            0, 2, 3, 1
        )
        orint_map_s19 = (
            L2Norm(self.conv_o19(score_featmaps_s19), dim=1)
            .permute(0, 2, 3, 1)
            .unsqueeze(-2)
        )
        score_featmaps_s19 = score_featmaps_s19 + score_featmaps_s17

        score_featmaps_s21 = F.leaky_relu(
            self.insnorm10(self.conv10(score_featmaps_s19))
        )
        score_map_s21 = self.insnorm_s21(self.conv_s21(score_featmaps_s21)).permute(
            0, 2, 3, 1
        )
        orint_map_s21 = (
            L2Norm(self.conv_o21(score_featmaps_s21), dim=1)
            .permute(0, 2, 3, 1)
            .unsqueeze(-2)
        )

        score_maps = torch.cat(
            (
                score_map_s3,
                score_map_s5,
                score_map_s7,
                score_map_s9,
                score_map_s11,
                score_map_s13,
                score_map_s15,
                score_map_s17,
                score_map_s19,
                score_map_s21,
            ),
            -1,
        )  # (B, H, W, C)

        orint_maps = torch.cat(
            (
                orint_map_s3,
                orint_map_s5,
                orint_map_s7,
                orint_map_s9,
                orint_map_s11,
                orint_map_s13,
                orint_map_s15,
                orint_map_s17,
                orint_map_s19,
                orint_map_s21,
            ),
            -2,
        )  # (B, H, W, 10, 2)

        # get each pixel probability in all scale
        scale_probs = soft_nms_3d(score_maps, ksize=15, com_strength=3.0)

        # get each pixel probability summary from all scale space and correspond scale value
        score_map, scale_map, orint_map = soft_max_and_argmax_1d(
            input=scale_probs,
            orint_maps=orint_maps,
            dim=-1,
            scale_list=self.scale_list,
            keepdim=True,
            com_strength1=self.score_com_strength,
            com_strength2=self.scale_com_strength,
        )

        return score_map, scale_map, orint_map

    @staticmethod
    def convO_init(m):
        if isinstance(m, nn.Conv2d):
            nn.init.zeros_(m.weight.data)
            try:
                nn.init.ones_(m.bias.data)
            except:
                pass