import torch
import torch.nn as nn
import torch.nn.functional as F
from mmcv.cnn import constant_init, normal_init, kaiming_init
import numpy as np
import cv2 
import mmcv
from mmcv.image import imread, imwrite
from mmcv.cnn import ConvModule
from mmcv.runner import BaseModule
import pdb

# norm_cfg = {
#     # format: layer_type: (abbreviation, module)
#     'BN': ('bn', nn.BatchNorm2d),
#     'SyncBN': ('bn', nn.SyncBatchNorm),
#     'GN': ('gn', nn.GroupNorm),
#     # and potentially 'SN'
# }
class NonLocal2DMultiLayerTripleInception3Groupx261(BaseModule):
    """Non-local module.

    See https://arxiv.org/abs/1711.07971 for details.

    Args:
        in_channels (int): Channels of the input feature map.
        reduction (int): Channel reduction ratio.
        use_scale (bool): Whether to scale pairwise_weight by 1/inter_channels.
        conv_cfg (dict): The config dict for convolution layers.
            (only applicable to conv_out)
        norm_cfg (dict): The config dict for normalization layers.
            (only applicable to conv_out)
        mode (str): Options are `embedded_gaussian` and `dot_product`.
    """
    def __init__(self,
                 in_channels1,
                 in_channels2,
                 in_channels3,
                 reduction=2,
                 use_scale=True,
                 use_group=False,
                 group_num = 4,                 
                 conv_cfg=None,
                 norm_cfg=None,
                 mode='embedded_gaussian'):
        super(NonLocal2DMultiLayerTripleInception3Groupx261, self).__init__()
        self.in_channels1 = in_channels1
        self.in_channels2 = in_channels2
        self.in_channels3 = in_channels3
        self.reduction = reduction
        self.use_scale = use_scale
        self.use_group = use_group
        self.group_num = group_num
        self.inter_channels = in_channels1 // reduction
        self.group_channels = self.inter_channels // self.group_num
        self.mode = mode
        self.softmaxvis_cnt = 0
        assert mode in ['embedded_gaussian', 'dot_product']

        # g, theta, phi are actually `nn.Conv2d`. Here we use ConvModule for
        # potential usage.
        self.dec_tar = ConvModule(
            self.in_channels1,
            self.inter_channels,
            kernel_size=1,
            stride=1,
            dilation=1,
            act_cfg=None)
        self.dec_ref_1 = ConvModule(
            self.in_channels2,
            self.inter_channels,
            kernel_size=1,
            stride=1,
            dilation=1,
            act_cfg=None)
        self.dec_ref_2 = ConvModule(
            self.in_channels3,
            self.inter_channels,
            kernel_size=1,
            stride=1,
            dilation=1,
            act_cfg=None)
        # self.w_tar_1 = ConvModule(
        #     2*self.inter_channels,
        #     2,
        #     kernel_size=1,
        #     stride=1,
        #     dilation=1,
        #     activation=None)
        # self.w_ref_1 = ConvModule(
        #     3*self.inter_channels,
        #     3,
        #     kernel_size=1,
        #     stride=1,
        #     dilation=1,
        #     activation=None)
        # self.w_ref_2 = ConvModule(
        #     3*self.inter_channels,
        #     3,
        #     kernel_size=1,
        #     stride=1,
        #     dilation=1,
        #     activation=None)
        # self.g = ConvModule(
        #     3*self.inter_channels,
        #     3*self.inter_channels,
        #     kernel_size=1,
        #     stride=1,
        #     dilation=1,
        #     groups = 3,
        #     activation=None)
        self.theta = ConvModule(
            self.in_channels1,
            self.inter_channels,
            kernel_size=1,
            stride=1,
            dilation=1,
            act_cfg=None)
        # self.phi = ConvModule(
        #     3*self.inter_channels,
        #     3*self.inter_channels,
        #     kernel_size=1,
        #     stride=1,
        #     dilation=1,
        #     groups = 3,
        #     activation=None)
        self.conv_out = ConvModule(
            self.inter_channels,
            self.in_channels1,
            kernel_size=1,
            stride=1,
            dilation=1,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=None)
        self.down_ref   = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        
        # self.max_tar_a  = nn.MaxPool2d(kernel_size=5, stride=1, padding=2)
        # self.max_tar_b  = nn.MaxPool2d(kernel_size=9, stride=1, padding=4)

        # self.max_ref1_a = nn.MaxPool2d(kernel_size=5, stride=1, padding=2)
        # self.max_ref1_b = nn.MaxPool2d(kernel_size=9, stride=1, padding=4)

        # self.max_ref2_a = nn.MaxPool2d(kernel_size=5, stride=1, padding=2)
        # self.max_ref2_b = nn.MaxPool2d(kernel_size=9, stride=1, padding=4)

        # self.max_tar_a  = nn.MaxPool2d(kernel_size=3, stride=1, dilation=3, padding=0)
        self.max_tar_b  = nn.MaxPool2d(kernel_size=5, stride=1, dilation=4, padding=0)

        # self.max_ref1_a = nn.MaxPool2d(kernel_size=3, stride=1, dilation=3, padding=0)
        self.max_ref1_b = nn.MaxPool2d(kernel_size=5, stride=1, dilation=4, padding=0)

        # self.max_ref2_a = nn.MaxPool2d(kernel_size=3, stride=1, dilation=3, padding=0)
        self.max_ref2_b = nn.MaxPool2d(kernel_size=5, stride=1, dilation=4, padding=0)

        self.init_weights()
        self.test_image_index = 0

    def normalize(self, arr):
        new_arr = (arr - arr.min()) * (1/(arr.max() - arr.min()) * 255)# + 50
        return new_arr

    def save_softmax(self, softweight, h, w, img, filename, norm = False, global_norm = False):
        n, _, img_h, img_w = img.shape
        #factor = img_h // h

        img = img.squeeze().permute(1,2,0).cpu().numpy().astype('float32')
        #pdb.set_trace()
        img = mmcv.imdenormalize(img, np.array([102.9801, 115.9465, 122.7717]).astype('float32'),
            np.array([1.0, 1.0, 1.0]).astype('float32'), False)
        #imwrite(img, filename + '_gt.jpg')
        
        softweight = softweight.view(n, h,w,-1).permute(0,3,1,2)
        softweight = F.upsample(softweight, (img_h,img_w)).squeeze(0).permute(1,2,0)
        #pdb.set_trace()
        #softweight = self.normalize(softweight)
        _, _, num = softweight.shape

        if global_norm:
            softweight = self.normalize(softweight)

        for i in range(num):
            weight = softweight[:,:,i]
            weight = weight.cpu().numpy().astype('float32')
            if norm:
                weight = self.normalize(weight)
            elif not global_norm:
                weight = 255 * weight
            gray = np.uint8(weight)
            gray = cv2.applyColorMap(gray, cv2.COLORMAP_JET)
            gray = 0.5 * img + 0.5 * gray
            imwrite(gray, filename + '_' + str(i) + '.jpg')


    def init_weights(self, std=0.01, zeros_init=True):
        for m in [self.dec_tar, self.dec_ref_1, self.dec_ref_2, self.theta]:        	
            normal_init(m.conv, std=std)
        if zeros_init:
            constant_init(self.conv_out.conv, 0)
        else:
            normal_init(self.conv_out.conv, std=std)

    def embedded_gaussian(self, theta_x, phi_x):
        # pairwise_weight: [N, HxW, HxW]
        pairwise_weight = torch.matmul(theta_x, phi_x)
        if self.use_scale:
            # theta_x.shape[-1] is `self.inter_channels`
            # pairwise_weight /= theta_x.shape[-1]**-0.5
            pairwise_weight /= theta_x.shape[-1]**0.5
        pairwise_weight = pairwise_weight.softmax(dim=-1)
        return pairwise_weight

    def dot_product(self, theta_x, phi_x):
        # pairwise_weight: [N, H_tar x W_tar, H_ref x W_ref]
        pairwise_weight = torch.matmul(theta_x, phi_x)
        pairwise_weight /= pairwise_weight.shape[-1]
        return pairwise_weight

    def forward(self, x_tar, x_ref_1, x_ref_2, img=None, use_relu=True):

        n_tar, _, h_tar, w_tar = x_tar.shape

        # x_tar_a = self.max_tar_a(x_tar)
        # x_tar_b = self.max_tar_b(x_tar)

        # x_ref_1_a = self.max_ref1_a(x_ref_1)
        # x_ref_1_b = self.max_ref1_b(x_ref_1)

        # x_ref_2   = self.down_ref(x_ref_2)
        # x_ref_2_a = self.max_ref2_a(x_ref_2)
        # x_ref_2_b = self.max_ref2_b(x_ref_2)

        # x_tar_a = self.max_tar_a(F.pad(x_tar, [3,3,3,3], "constant", 0))
        x_tar_b = self.max_tar_b(F.pad(x_tar, [8,8,8,8], "constant", 0))

        # x_ref_1   = x_ref_1
        # x_ref_1_a = self.max_ref1_a(F.pad(x_ref_1, [3,3,3,3], "constant", 0))
        x_ref_1_b = self.max_ref1_b(F.pad(x_ref_1, [8,8,8,8], "constant", 0))

        x_ref_2   = self.down_ref(x_ref_2)
        # x_ref_2_a = self.max_ref2_a(F.pad(x_ref_2, [3,3,3,3], "constant", 0))
        x_ref_2_b = self.max_ref2_b(F.pad(x_ref_2, [8,8,8,8], "constant", 0))


        # x_tar_a = x_tar_a.view(n_tar, self.in_channels1, h_tar*w_tar, 1) 
        x_tar_b = x_tar_b.view(n_tar, self.in_channels1, h_tar*w_tar, 1) 
        # x_tar: [N, C_in1, H_tar*W_tar, 1]

        # x_ref_1     = x_ref_1.view(n_tar, self.in_channels2, h_tar*w_tar, 1) 
        # x_ref_1_a = x_ref_1_a.view(n_tar, self.in_channels2, h_tar*w_tar, 1) 
        x_ref_1_b = x_ref_1_b.view(n_tar, self.in_channels2, h_tar*w_tar, 1) 
        # x_ref_1: [N, C_in2, H_tar*W_tar, 1]

        # x_ref_2     = x_ref_2.view(n_tar, self.in_channels3, h_tar*w_tar, 1) 
        # x_ref_2_a = x_ref_2_a.view(n_tar, self.in_channels3, h_tar*w_tar, 1) 
        x_ref_2_b = x_ref_2_b.view(n_tar, self.in_channels3, h_tar*w_tar, 1) 
        # x_ref_2: [N, C_in3, H_tar*W_tar, 1]



        # x_tar_cat = torch.cat((x_tar_a, x_tar_b), 1)
        # # x_tar_cat: [N, 2*C_in1, H_tar*W_tar, 1]
        # dec_tar = self.dec_tar(x_tar_cat).permute(0, 2, 3, 1).view(n_tar, h_tar*w_tar, 2, self.inter_channels)
        # # dec_tar: [N, 2*C_inter, H_tar*W_tar, 1] ->  [N, H_tar*W_tar, 2, C_inter]
        # x_tar_cat = dec_tar.reshape(n_tar, h_tar*w_tar, -1).permute(0, 2, 1).unsqueeze(dim=3)
        # # x_tar_cat: [N, H_tar*W_tar, 2*C_inter] -> [N, 2*C_inter, H_tar*W_tar, 1]
        # w_tar_1 = self.w_tar_1(x_tar_cat).softmax(dim=1).permute(0, 2, 3, 1)
        # # w_tar_1: [N, 2, H_tar*W_tar, 1] -> [N, H_tar*W_tar, 1, 2]
        # x_tar_1 = torch.matmul(w_tar_1, dec_tar).permute(0, 3, 1, 2)
        # # x_tar_1: [N, H_tar*W_tar, 1, C_inter] -> [N, C_inter, H_tar*W_tar, 1]
        x_tar_1 = self.dec_tar(x_tar_b)
        # x_tar_1: [N, C_inter, H_tar*W_tar, 1]


        # x_ref_1_cat = torch.cat((x_ref_1, x_ref_1_a, x_ref_1_b), 1)
        # # x_ref_1_cat: [N, 3*C_in2, H_tar*W_tar, 1]
        # dec_ref_1 = self.dec_ref_1(x_ref_1_cat).permute(0, 2, 3, 1).view(n_tar, h_tar*w_tar, 3, self.inter_channels)
        # # dec_ref_1: [N, 3*C_inter, H_tar*W_tar, 1] -> [N, H_tar*W_tar, 3, C_inter]
        # x_ref_1_cat = dec_ref_1.reshape(n_tar, h_tar*w_tar, -1).permute(0, 2, 1).unsqueeze(dim=3)
        # # x_ref_1_cat: [N, 3*C_inter, H_tar*W_tar, 1]
        # w_ref_1 = self.w_ref_1(x_ref_1_cat).softmax(dim=1).permute(0, 2, 3, 1)
        # # w_ref_1: [N, 3, H_tar*W_tar, 1] -> [N, H_tar*W_tar, 1, 3]
        # x_ref_1 = torch.matmul(w_ref_1, dec_ref_1).permute(0, 3, 1, 2)
        # # x_ref_1: [N, H_tar*W_tar, 1, C_inter] -> [N, C_inter, H_tar*W_tar, 1]
        x_ref_1 = self.dec_ref_1(x_ref_1_b)
        # x_ref_1: [N, C_inter, H_tar*W_tar, 1]


        # x_ref_2_cat = torch.cat((x_ref_2, x_ref_2_a, x_ref_2_b), 1)
        # # x_ref_2_cat: [N, 3*C_in2, H_tar*W_tar, 1]
        # dec_ref_2 = self.dec_ref_2(x_ref_2_cat).permute(0, 2, 3, 1).view(n_tar, h_tar*w_tar, 3, self.inter_channels)
        # # dec_ref_2: [N, 3*C_inter, H_tar*W_tar, 1] -> [N, H_tar*W_tar, 3, C_inter]
        # x_ref_2_cat = dec_ref_2.reshape(n_tar, h_tar*w_tar, -1).permute(0, 2, 1).unsqueeze(dim=3)
        # # x_ref_2_cat: [N, 3*C_inter, H_tar*W_tar, 1]
        # w_ref_2 = self.w_ref_2(x_ref_2_cat).softmax(dim=1).permute(0, 2, 3, 1)
        # # w_ref_2: [N, 3, H_tar*W_tar, 1] -> [N, H_tar*W_tar, 1, 3]
        # x_ref_2 = torch.matmul(w_ref_2, dec_ref_2).permute(0, 3, 1, 2)
        # # x_ref_2: [N, H_tar*W_tar, 1, C_inter] -> [N, C_inter, H_tar*W_tar, 1]
        x_ref_2 = self.dec_ref_2(x_ref_2_b)
        # x_ref_2: [N, C_inter, H_tar*W_tar, 1]


        x_ref = torch.cat((x_tar_1, x_ref_1, x_ref_2), 1)
        # x_ref: [N, 3*C_inter, H_tar*W_tar, 1]

        # g_x = self.g(x_ref)
        g_x = x_ref
        # g_x: [N, 3*C_inter, H_tar*W_tar, 1]
        g_x = g_x.squeeze(dim=3).permute(0, 2, 1).view(n_tar, h_tar*w_tar, 3, self.inter_channels)
        # g_x: [N, H_tar*W_tar, 3, C_inter]
        if self.use_group:
            g_x = g_x.view(n_tar, h_tar*w_tar, 3, self.group_num, self.group_channels)
            # g_x: [N, H_tar*W_tar, 3, group_num, C_group]
            g_x = g_x.permute(0, 1, 3, 2, 4)
            # g_x: [N, H_tar*W_tar, group_num, 3, C_group]
            g_x = g_x.reshape(n_tar, h_tar*w_tar*self.group_num, 3, self.group_channels)
            # g_x: [N, H_tar*W_tar*group_num, 3, C_group]        


        theta_x = self.theta(x_tar)
        # theta_x: [N, C_inter, H_tar, W_tar]
        theta_x = theta_x.view(n_tar, self.inter_channels, -1).permute(0, 2, 1)
        # theta_x: [N, H_tarxW_tar, C_inter]
        theta_x = theta_x[:, :, None, :]
        # theta_x: [N, H_tarxW_tar, 1, C_inter]
        if self.use_group:
            theta_x = theta_x.view(n_tar, h_tar*w_tar, 1, self.group_num, self.group_channels)
            # theta_x: [N, H_tar*W_tar, 1, group_num, C_group]
            theta_x = theta_x.permute(0, 1, 3, 2, 4)
            # theta_x: [N, H_tar*W_tar, group_num, 1, C_group]
            theta_x = theta_x.reshape(n_tar, h_tar*w_tar*self.group_num, 1, self.group_channels)
            # theta_x: [N, H_tar*W_tar*group_num, 1, C_group]    


        # phi_x = self.phi(x_ref)
        phi_x = x_ref
        # phi_x: [N, 3*C_inter, H_tar*W_tar, 1]
        phi_x = phi_x.squeeze(dim=3).permute(0, 2, 1).view(n_tar, h_tar*w_tar, 3, self.inter_channels)
        # phi_x: [N, H_tar*W_tar, 3, C_inter]
        phi_x = phi_x.permute(0, 1, 3, 2)
        # phi_x: [N, H_tarxW_tar, C_inter, 3]
        if self.use_group:
            phi_x = phi_x.reshape(n_tar, h_tar*w_tar*self.group_num, self.group_channels, 3)
            # phi_x: [N, H_tar*W_tar*group_num, C_group, 3]


        pairwise_func = getattr(self, self.mode)
        pairwise_weight = pairwise_func(theta_x, phi_x)
        #self.save_softmax(pairwise_weight, h_tar, w_tar, img, 'result_nlmb/' + str(self.test_image_index).zfill(6)  + '/softmax')
        # self.save_softmax(pairwise_weight, h_tar, w_tar, 'softmaxvis/softmax')
        # pairwise_weight = F.relu(pairwise_weight)
        # pairwise_weight: [N, H_tarxW_tar(xgroup_num), 1, 3]           

        y = torch.matmul(pairwise_weight, g_x)
        # y: [N, H_tarxW_tar(xgroup_num), 1, C_inter(C_group)]          
        y = y.squeeze(dim=2)
        # y: [N, H_tarxW_tar(xgroup_num), C_inter(C_group)]
        if self.use_group:
            y = y.view(n_tar, h_tar*w_tar, self.inter_channels)
            # y: [N, H_tarxW_tar, C_inter]

        y = y.permute(0, 2, 1).reshape(n_tar, self.inter_channels, h_tar, w_tar)
        # y: [N, C_inter, H_tar, W_tar]
        y = self.conv_out(y)
        # y: [N, C_tar, H_tar, W_tar]

        output = x_tar + y

        #self.save_softmax(x_tar.mean(dim=1).unsqueeze(3), h_tar, w_tar, img, 
        #    'result_nlmb/' + str(self.test_image_index).zfill(6)  + '/tar_avg', True)
        #x_tar_max, _ = x_tar.max(dim=1)
        #self.save_softmax(x_tar_max.unsqueeze(3), h_tar, w_tar, img, 
        #    'result_nlmb/' + str(self.test_image_index).zfill(6)  + '/tar_max', True)

        #self.save_softmax(F.relu(y).mean(dim=1).unsqueeze(3), h_tar, w_tar, img, 
        #    'result_nlmb/' + str(self.test_image_index).zfill(6)  + '/y_relu_avg', True)
        #y_max, _ = F.relu(y).max(dim=1)
        #self.save_softmax(y_max.unsqueeze(3), h_tar, w_tar, img, 
        #    'result_nlmb/' + str(self.test_image_index).zfill(6)  + '/y_relu_max', True)

        # output = self.bn2(output)
        if use_relu:
            output = F.relu(output)

        #feature_map = torch.cat((x_tar.mean(dim=1).unsqueeze(3), output.mean(dim=1).unsqueeze(3)), dim=3)
        
        #self.save_softmax(feature_map, h_tar, w_tar, img, 
        #    'result_mask_feature/' + str(self.test_image_index).zfill(6)  + '/feature_map', norm=True)
        #self.save_softmax(output.permute(0,2,3,1), h_tar, w_tar, img, 
        #    'result_mask_feature/output', norm=True)
        #self.save_softmax(x_tar.permute(0,2,3,1), h_tar, w_tar, img, 
        #    'result_mask_feature/x_tar', norm=True)
        #self.save_softmax(output.mean(dim=1).unsqueeze(3), h_tar, w_tar, img, 
        #    'result_nlmb/' + str(self.test_image_index).zfill(6)  + '/output_avg', True)
        #output_max, _ = output.max(dim=1)
        #self.save_softmax(output_max.unsqueeze(3), h_tar, w_tar, img, 
        #    'result_nlmb/' + str(self.test_image_index).zfill(6)  + '/output_max', True)

        #self.test_image_index = self.test_image_index + 1

        return output
