"""
Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
"""

import torch
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
from models.networks.base_network import BaseNetwork
from models.networks.normalization import get_nonspade_norm_layer
import util.util as util

upsampling1 = nn.Upsample((129,129))
upsampling2 = nn.Upsample((65,65))


class MultiscaleDiscriminator(BaseNetwork):
    @staticmethod
    def modify_commandline_options(parser, is_train):
        parser.add_argument('--netD_subarch', type=str, default='n_layer',
                            help='architecture of each discriminator')
        parser.add_argument('--num_D', type=int, default=2,
                            help='number of discriminators to be used in multiscale')
        opt, _ = parser.parse_known_args()

        # define properties of each discriminator of the multiscale discriminator
        subnetD = util.find_class_in_module(opt.netD_subarch + 'discriminator',
                                            'models.networks.discriminator')
        subnetD.modify_commandline_options(parser, is_train)

        return parser

    def __init__(self, opt):
        super().__init__()
        self.opt = opt

        for i in range(opt.num_D):
            subnetD = self.create_single_discriminator(opt)
            self.add_module('discriminator_%d' % i, subnetD)

    def create_single_discriminator(self, opt):
        subarch = opt.netD_subarch
        if subarch == 'n_layer':
            netD = NLayerDiscriminator(opt)
        else:
            raise ValueError('unrecognized discriminator subarchitecture %s' % subarch)
        return netD

    def downsample(self, input):
        return F.avg_pool2d(input, kernel_size=3,
                            stride=2, padding=[1, 1],
                            count_include_pad=False)

    # Returns list of lists of discriminator outputs.
    # The final result is of size opt.num_D x opt.n_layers_D
    def forward(self, input, set_1_2=None, D_feature_SEAN=None):
        result = []
        get_intermediate_features = not self.opt.no_ganFeat_loss
        for name, D in self.named_children():
            out = D(input, set_1_2=set_1_2, D_feature_SEAN=D_feature_SEAN, name=name)
            if not get_intermediate_features:
                out = [out]
            result.append(out)
            input = self.downsample(input)

        return result


# Defines the PatchGAN discriminator with the specified arguments.
class NLayerDiscriminator(BaseNetwork):
    @staticmethod
    def modify_commandline_options(parser, is_train):
        parser.add_argument('--n_layers_D', type=int, default=3,
                            help='# layers in each discriminator')
        return parser

    def __init__(self, opt):
        super().__init__()
        self.opt = opt

        kw = 4
        padw = int(np.ceil((kw - 1.0) / 2))
        nf = opt.ndf
        input_nc = self.compute_D_input_nc(opt)

        norm_layer = get_nonspade_norm_layer(opt, opt.norm_D)
        sequence = [[nn.Conv2d(input_nc, nf, kernel_size=kw, stride=2, padding=padw),
                     nn.LeakyReLU(0.2, False)]]

        for n in range(1, opt.n_layers_D):
            nf_prev = nf
            nf = min(nf * 2, 512)
            sequence += [[norm_layer(nn.Conv2d(nf_prev, nf, kernel_size=kw,
                                               stride=2, padding=padw)),
                          nn.LeakyReLU(0.2, False)
                          ]]

        sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)]]

        # We divide the layers into groups to extract intermediate layer outputs
        for n in range(len(sequence)):
            self.add_module('model' + str(n), nn.Sequential(*sequence[n]))

    def compute_D_input_nc(self, opt):
        input_nc = opt.label_nc + opt.output_nc
        if opt.contain_dontcare_label:
            input_nc += 1
        if not opt.no_instance:
            input_nc += 1
        return input_nc

    def forward(self, input, set_1_2=None, D_feature_SEAN=None, name=None):
        results = [input]
        #set1, set2= set_1_2
        #gen_in1, Fake_img_1_1, G_feature_1_1, D_feature_1_1, Fake_img_1_2, G_feature_1_2, D_feature_1_2, Fake_img_1_mix_1_2, G_feature_1_mix_1_2, D_feature_1_mix_1_2 = set1  
        #gen_in2, Fake_img_2_1, G_feature_2_1, D_feature_2_1, Fake_img_2_2, G_feature_2_2, D_feature_2_2, Fake_img_2_mix_1_2, G_feature_2_mix_1_2, D_feature_2_mix_1_2 = set2 

        #gen_in1, Fake_img_1_1, G_feature_1_1, D_feature_1_1, Fake_img_1_2, G_feature_1_2, D_feature_1_2, Fake_img_1_mix_1_2, G_feature_1_mix_1_2, D_feature_1_mix_1_2 = set1  
        #D_feature_real=set_1_2[0][3]# since we use G_feature_1_1 in generator 
        #r1,r2 = torch.split(D_feature_real[128],64,dim=1)
        #r_1_2 = (r1+r2)/2.0
        #f1,f2 = torch.split(D_feature_SEAN[128],64,dim=1)
        #f_1_2 = (f1+f2)/2.0
        #intermediate_output = torch.cat([f_1_2, r_1_2], dim=0)
        #if name=='discriminator_0':
        #    intermediate_output=upsampling1(intermediate_output)
        #else:
        #    intermediate_output=upsampling2(intermediate_output)
        #results = [intermediate_output]
        for submodel_ind, submodel in enumerate(self.children()):
            intermediate_output = submodel(results[-1])
            results.append(intermediate_output)

        #del set1 
        #del set2 
        get_intermediate_features = not self.opt.no_ganFeat_loss
        if get_intermediate_features:
            return results[1:]
        else:
            return results[-1]
