# -*- coding: utf-8 -*-
# @Time    : 2018-9-27 15:39
# @Author  : xylon

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

from utils.image_utils import filter_border, nms, topk_map, get_gauss_filter_weight


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

        self.score_com_strength = score_com_strength
        self.scale_com_strength = scale_com_strength
        self.NMS_THRESH = nms_thresh
        self.NMS_KSIZE = nms_ksize
        self.TOPK = topk
        self.GAUSSIAN_KSIZE = gauss_ksize
        self.GAUSSIAN_SIGMA = gauss_sigma

        self.conv1 = nn.Conv2d(
            in_channels=1,
            out_channels=16,
            kernel_size=ksize,
            stride=1,
            padding=padding,
            dilation=dilation,
        )  # 3 RF
        self.insnorm1 = nn.InstanceNorm2d(16, affine=True)
        self.conv_s3 = nn.Conv2d(
            in_channels=16, out_channels=1, kernel_size=1, stride=1, padding=0
        )
        self.insnorm_s3 = nn.InstanceNorm2d(1, affine=True)

        self.conv2 = nn.Conv2d(
            in_channels=16,
            out_channels=16,
            kernel_size=ksize,
            stride=1,
            padding=padding,
            dilation=dilation,
        )  # 5 RF
        self.insnorm2 = nn.InstanceNorm2d(16, affine=True)
        self.conv_s5 = nn.Conv2d(
            in_channels=16, out_channels=1, kernel_size=1, stride=1, padding=0
        )
        self.insnorm_s5 = nn.InstanceNorm2d(1, affine=True)

        self.conv3 = nn.Conv2d(
            in_channels=16,
            out_channels=16,
            kernel_size=ksize,
            stride=1,
            padding=padding,
            dilation=dilation,
        )  # 7 RF
        self.insnorm3 = nn.InstanceNorm2d(16, affine=True)
        self.conv_s7 = nn.Conv2d(
            in_channels=16, out_channels=1, kernel_size=1, stride=1, padding=0
        )
        self.insnorm_s7 = nn.InstanceNorm2d(1, affine=True)

        self.conv4 = nn.Conv2d(
            in_channels=16,
            out_channels=16,
            kernel_size=ksize,
            stride=1,
            padding=padding,
            dilation=dilation,
        )  # 9 RF
        self.insnorm4 = nn.InstanceNorm2d(16, affine=True)
        self.conv_s9 = nn.Conv2d(
            in_channels=16, out_channels=1, kernel_size=1, stride=1, padding=0
        )
        self.insnorm_s9 = nn.InstanceNorm2d(1, affine=True)

        self.conv5 = nn.Conv2d(
            in_channels=16,
            out_channels=16,
            kernel_size=ksize,
            stride=1,
            padding=padding,
            dilation=dilation,
        )  # 11 RF
        self.insnorm5 = nn.InstanceNorm2d(16, affine=True)
        self.conv_s11 = nn.Conv2d(
            in_channels=16, out_channels=1, kernel_size=1, stride=1, padding=0
        )
        self.insnorm_s11 = nn.InstanceNorm2d(1, affine=True)

        self.conv6 = nn.Conv2d(
            in_channels=16,
            out_channels=16,
            kernel_size=ksize,
            stride=1,
            padding=padding,
            dilation=dilation,
        )  # 13 RF
        self.insnorm6 = nn.InstanceNorm2d(16, affine=True)
        self.conv_s13 = nn.Conv2d(
            in_channels=16, out_channels=1, kernel_size=1, stride=1, padding=0
        )
        self.insnorm_s13 = nn.InstanceNorm2d(1, affine=True)

        self.conv7 = nn.Conv2d(
            in_channels=16,
            out_channels=16,
            kernel_size=ksize,
            stride=1,
            padding=padding,
            dilation=dilation,
        )  # 15 RF
        self.insnorm7 = nn.InstanceNorm2d(16, affine=True)
        self.conv_s15 = nn.Conv2d(
            in_channels=16, out_channels=1, kernel_size=1, stride=1, padding=0
        )
        self.insnorm_s15 = nn.InstanceNorm2d(1, affine=True)

        self.conv8 = nn.Conv2d(
            in_channels=16,
            out_channels=16,
            kernel_size=ksize,
            stride=1,
            padding=padding,
            dilation=dilation,
        )  # 17 RF
        self.insnorm8 = nn.InstanceNorm2d(16, affine=True)
        self.conv_s17 = nn.Conv2d(
            in_channels=16, out_channels=1, kernel_size=1, stride=1, padding=0
        )
        self.insnorm_s17 = nn.InstanceNorm2d(1, affine=True)

        self.conv9 = nn.Conv2d(
            in_channels=16,
            out_channels=16,
            kernel_size=ksize,
            stride=1,
            padding=padding,
            dilation=dilation,
        )  # 19 RF
        self.insnorm9 = nn.InstanceNorm2d(16, affine=True)
        self.conv_s19 = nn.Conv2d(
            in_channels=16, out_channels=1, kernel_size=1, stride=1, padding=0
        )
        self.insnorm_s19 = nn.InstanceNorm2d(1, affine=True)

        self.conv10 = nn.Conv2d(
            in_channels=16,
            out_channels=16,
            kernel_size=ksize,
            stride=1,
            padding=padding,
            dilation=dilation,
        )  # 21 RF
        self.insnorm10 = nn.InstanceNorm2d(16, affine=True)
        self.conv_s21 = nn.Conv2d(
            in_channels=16, out_channels=1, kernel_size=1, stride=1, padding=0
        )
        self.insnorm_s21 = nn.InstanceNorm2d(1, affine=True)

        self.scale_list = torch.tensor(scale_list)

    def forward(self, **kwargs):
        pass

    def process(self, im1w_score):
        """
        nms(n), topk(t), gaussian kernel(g) operation
        :param im1w_score: warped score map
        :return: processed score map, topk mask, topk value
        """
        im1w_score = filter_border(im1w_score)

        # apply nms to im1w_score
        nms_mask = nms(im1w_score, thresh=self.NMS_THRESH, ksize=self.NMS_KSIZE)
        im1w_score = im1w_score * nms_mask
        topk_value = im1w_score

        # apply topk to im1w_score
        topk_mask = topk_map(im1w_score, self.TOPK)
        im1w_score = topk_mask.to(torch.float) * im1w_score

        # apply gaussian kernel to im1w_score
        psf = im1w_score.new_tensor(
            get_gauss_filter_weight(self.GAUSSIAN_KSIZE, self.GAUSSIAN_SIGMA)[
                None, None, :, :
            ]
        ).to(im1w_score.device)
        im1w_score = F.conv2d(
            input=im1w_score.permute(0, 3, 1, 2),
            weight=psf,
            stride=1,
            padding=self.GAUSSIAN_KSIZE // 2,
        ).permute(
            0, 2, 3, 1
        )  # (B, H, W, 1)

        """
        apply tf.clamp to make sure all value in im1w_score isn't greater than 1
        but this won't happend in correct way
        """
        im1w_score = im1w_score.clamp(min=0.0, max=1.0)

        return im1w_score, topk_mask, topk_value

    # @staticmethod
    # def loss(left_score, im1gt_score, im1visible_mask):
    #     im1_score = left_score

    #     l2_element_diff = (im1_score - im1gt_score) ** 2
    #     # visualization numbers
    #     Nvi = torch.clamp(im1visible_mask.sum(dim=(3, 2, 1)), min=2.0)
    #     loss = (
    #         torch.sum(l2_element_diff * im1visible_mask, dim=(3, 2, 1)) / (Nvi + 1e-8)
    #     ).mean()

    #     return loss

    @staticmethod
    def weights_init(m):
        if isinstance(m, nn.Conv2d):
            nn.init.xavier_uniform_(
                m.weight.data, gain=nn.init.calculate_gain("leaky_relu")
            )
            try:
                nn.init.xavier_uniform_(m.bias.data)
            except:
                pass
