from unittest import skip
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import init
from torch.nn.utils import spectral_norm
from torch.optim import lr_scheduler
from torch.autograd import Variable
import functools
import numpy as np
import math

class Identity(nn.Module):
    def forward(self, x):
        return x

def get_non_linearity(layer_type='relu'):
    if layer_type == 'relu':
        nl_layer = functools.partial(nn.ReLU, inplace=True)
    elif layer_type == 'lrelu':
        nl_layer = functools.partial(
            nn.LeakyReLU, negative_slope=0.2, inplace=True)
    elif layer_type == 'elu':
        nl_layer = functools.partial(nn.ELU, inplace=True)
    else:
        raise NotImplementedError(
            'nonlinearity activitation [%s] is not found' % layer_type)
    return nl_layer

def get_norm_layer(norm_type='instance'):
    """Return a normalization layer

    Parameters:
        norm_type (str) -- the name of the normalization layer: batch | instance | none

    For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).
    For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.
    """
    if norm_type == 'batch':
        norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
    elif norm_type == 'instance':
        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
    elif norm_type == 'none':
        def norm_layer(x): return Identity()
    else:
        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
    return norm_layer


def get_scheduler(optimizer, opt):
    """Return a learning rate scheduler

    Parameters:
        optimizer          -- the optimizer of the network
        opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions．　
                              opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine

    For 'linear', we keep the same learning rate for the first <opt.n_epochs> epochs
    and linearly decay the rate to zero over the next <opt.n_epochs_decay> epochs.
    For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers.
    See https://pytorch.org/docs/stable/optim.html for more details.
    """
    if opt.lr_policy == 'linear':
        def lambda_rule(epoch):
            lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.n_epochs) / float(opt.n_epochs_decay + 1)
            return lr_l
        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
    elif opt.lr_policy == 'step':
        scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)
    elif opt.lr_policy == 'plateau':
        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
    elif opt.lr_policy == 'cosine':
        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.n_epochs, eta_min=0)
    else:
        return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)
    return scheduler


def init_weights(net, init_type='normal', init_gain=0.02):
    """Initialize network weights.

    Parameters:
        net (network)   -- network to be initialized
        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.

    We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
    work better for some applications. Feel free to try yourself.
    """
    def init_func(m):  # define the initialization function
        classname = m.__class__.__name__
        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
            if init_type == 'normal':
                init.normal_(m.weight.data, 0.0, init_gain)
            elif init_type == 'xavier':
                init.xavier_normal_(m.weight.data, gain=init_gain)
            elif init_type == 'kaiming':
                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
            elif init_type == 'orthogonal':
                init.orthogonal_(m.weight.data, gain=init_gain)
            else:
                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
            if hasattr(m, 'bias') and m.bias is not None:
                init.constant_(m.bias.data, 0.0)
        elif classname.find('BatchNorm2d') != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
            init.normal_(m.weight.data, 1.0, init_gain)
            init.constant_(m.bias.data, 0.0)

    print('initialize network with %s' % init_type)
    net.apply(init_func)  # apply the initialization function <init_func>


def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
    """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights
    Parameters:
        net (network)      -- the network to be initialized
        init_type (str)    -- the name of an initialization method: normal | xavier | kaiming | orthogonal
        gain (float)       -- scaling factor for normal, xavier and orthogonal.
        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2

    Return an initialized network.
    """
    if len(gpu_ids) > 0:
        assert(torch.cuda.is_available())
        net.to(gpu_ids[0])
        net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs
    init_weights(net, init_type, init_gain=init_gain)
    return net


def define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[]):
    """Create a generator

    Parameters:
        input_nc (int) -- the number of channels in input images
        output_nc (int) -- the number of channels in output images
        ngf (int) -- the number of filters in the last conv layer
        netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128
        norm (str) -- the name of normalization layers used in the network: batch | instance | none
        use_dropout (bool) -- if use dropout layers.
        init_type (str)    -- the name of our initialization method.
        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.
        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2

    Returns a generator

    Our current implementation provides two types of generators:
        U-Net: [unet_128] (for 128x128 input images) and [unet_256] (for 256x256 input images)
        The original U-Net paper: https://arxiv.org/abs/1505.04597

        Resnet-based generator: [resnet_6blocks] (with 6 Resnet blocks) and [resnet_9blocks] (with 9 Resnet blocks)
        Resnet-based generator consists of several Resnet blocks between a few downsampling/upsampling operations.
        We adapt Torch code from Justin Johnson's neural style transfer project (https://github.com/jcjohnson/fast-neural-style).


    The generator has been initialized by <init_net>. It uses RELU for non-linearity.
    """
    net = None
    norm_layer = get_norm_layer(norm_type=norm)

    if netG=='MFNET_G':
        net = MFNet_Generator(ngf=ngf, use_dropout=use_dropout, n_blocks=6)
    else:
        raise NotImplementedError('Generator model name [%s] is not recognized' % netG)
    return init_net(net, init_type, init_gain, gpu_ids)


def define_D(input_nc, ndf, netD, n_layers_D=3, norm='batch', init_type='normal', init_gain=0.02, gpu_ids=[], use_spectral_norm=False):
    """Create a discriminator

    Parameters:
        input_nc (int)     -- the number of channels in input images
        ndf (int)          -- the number of filters in the first conv layer
        netD (str)         -- the architecture's name: basic | n_layers | pixel
        n_layers_D (int)   -- the number of conv layers in the discriminator; effective when netD=='n_layers'
        norm (str)         -- the type of normalization layers used in the network.
        init_type (str)    -- the name of the initialization method.
        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.
        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2

    Returns a discriminator

    Our current implementation provides three types of discriminators:
        [basic]: 'PatchGAN' classifier described in the original pix2pix paper.
        It can classify whether 70×70 overlapping patches are real or fake.
        Such a patch-level discriminator architecture has fewer parameters
        than a full-image discriminator and can work on arbitrarily-sized images
        in a fully convolutional fashion.

        [n_layers]: With this mode, you can specify the number of conv layers in the discriminator
        with the parameter <n_layers_D> (default=3 as used in [basic] (PatchGAN).)

        [pixel]: 1x1 PixelGAN discriminator can classify whether a pixel is real or not.
        It encourages greater color diversity but has no effect on spatial statistics.

    The discriminator has been initialized by <init_net>. It uses Leakly RELU for non-linearity.
    """
    net = None
    norm_layer = get_norm_layer(norm_type=norm)

    if netD == 'basic':  # default PatchGAN classifier Receptive Field = 70
        net = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer)
    elif netD == 'basic_64':  # default PatchGAN classifier Receptive Field = 34
        if use_spectral_norm:
            net = NLayerDiscriminatorS(input_nc, ndf, n_layers=2, norm_layer=norm_layer)
        else:
            net = NLayerDiscriminator(input_nc, ndf, n_layers=2, norm_layer=norm_layer)
    else:
        raise NotImplementedError('Discriminator model name [%s] is not recognized' % netD)
    return init_net(net, init_type, init_gain, gpu_ids)


##############################################################################
# Classes
##############################################################################
class GANLoss(nn.Module):
    """Define different GAN objectives.

    The GANLoss class abstracts away the need to create the target label tensor
    that has the same size as the input.
    """

    def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):
        """ Initialize the GANLoss class.

        Parameters:
            gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp.
            target_real_label (bool) - - label for a real image
            target_fake_label (bool) - - label of a fake image

        Note: Do not use sigmoid as the last layer of Discriminator.
        LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss.
        """
        super(GANLoss, self).__init__()
        self.register_buffer('real_label', torch.tensor(target_real_label))
        self.register_buffer('fake_label', torch.tensor(target_fake_label))
        self.gan_mode = gan_mode
        if gan_mode == 'lsgan':
            self.loss = nn.MSELoss()
        elif gan_mode == 'vanilla':
            self.loss = nn.BCEWithLogitsLoss()
        elif gan_mode in ['wgangp', 'hinge']:
            self.loss = None
        else:
            raise NotImplementedError('gan mode %s not implemented' % gan_mode)

    def get_target_tensor(self, prediction, target_is_real):
        """Create label tensors with the same size as the input.

        Parameters:
            prediction (tensor) - - tpyically the prediction from a discriminator
            target_is_real (bool) - - if the ground truth label is for real images or fake images

        Returns:
            A label tensor filled with ground truth label, and with the size of the input
        """

        if target_is_real:
            target_tensor = self.real_label
        else:
            target_tensor = self.fake_label
        return target_tensor.expand_as(prediction)

    def __call__(self, prediction, target_is_real, train_gen=False):
        """Calculate loss given Discriminator's output and grount truth labels.

        Parameters:
            prediction (tensor) - - tpyically the prediction output from a discriminator
            target_is_real (bool) - - if the ground truth label is for real images or fake images

        Returns:
            the calculated loss.
        """
        if self.gan_mode in ['lsgan', 'vanilla']:
            target_tensor = self.get_target_tensor(prediction, target_is_real)
            loss = self.loss(prediction, target_tensor)
        elif self.gan_mode == 'wgangp':
            if target_is_real:
                loss = -prediction.mean()
            else:
                loss = prediction.mean()
        elif self.gan_mode == 'hinge':
            if train_gen:
                loss = -prediction.mean()
            else:
                if target_is_real:
                    loss = torch.nn.ReLU()(1.0 - prediction).mean()
                else:
                    loss =  torch.nn.ReLU()(1.0 + prediction).mean()
        return loss


def cal_gradient_penalty(netD, real_data, fake_data, device, type='mixed', constant=1.0, lambda_gp=10.0):
    """Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028

    Arguments:
        netD (network)              -- discriminator network
        real_data (tensor array)    -- real images
        fake_data (tensor array)    -- generated images from the generator
        device (str)                -- GPU / CPU: from torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')
        type (str)                  -- if we mix real and fake data or not [real | fake | mixed].
        constant (float)            -- the constant used in formula ( | |gradient||_2 - constant)^2
        lambda_gp (float)           -- weight for this loss

    Returns the gradient penalty loss
    """
    if lambda_gp > 0.0:
        if type == 'real':   # either use real images, fake images, or a linear interpolation of two.
            interpolatesv = real_data
        elif type == 'fake':
            interpolatesv = fake_data
        elif type == 'mixed':
            alpha = torch.rand(real_data.shape[0], 1, device=device)
            alpha = alpha.expand(real_data.shape[0], real_data.nelement() // real_data.shape[0]).contiguous().view(*real_data.shape)
            interpolatesv = alpha * real_data + ((1 - alpha) * fake_data)
        else:
            raise NotImplementedError('{} not implemented'.format(type))
        interpolatesv.requires_grad_(True)
        disc_interpolates = netD(interpolatesv)
        gradients = torch.autograd.grad(outputs=disc_interpolates, inputs=interpolatesv,
                                        grad_outputs=torch.ones(disc_interpolates.size()).to(device),
                                        create_graph=True, retain_graph=True, only_inputs=True)
        gradients = gradients[0].view(real_data.size(0), -1)  # flat the data
        gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) - constant) ** 2).mean() * lambda_gp        # added eps
        return gradient_penalty, gradients
    else:
        return 0.0, None
            
class ResnetGenerator(nn.Module):
    """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.

    We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)
    """

    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
        """Construct a Resnet-based generator

        Parameters:
            input_nc (int)      -- the number of channels in input images
            output_nc (int)     -- the number of channels in output images
            ngf (int)           -- the number of filters in the last conv layer
            norm_layer          -- normalization layer
            use_dropout (bool)  -- if use dropout layers
            n_blocks (int)      -- the number of ResNet blocks
            padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero
        """
        assert(n_blocks >= 0)
        super(ResnetGenerator, self).__init__()
        if type(norm_layer) == functools.partial:
            use_bias = norm_layer.func == nn.InstanceNorm2d
        else:
            use_bias = norm_layer == nn.InstanceNorm2d

        model = [nn.ReflectionPad2d(3),
                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
                 norm_layer(ngf),
                 nn.ReLU(True)]

        n_downsampling = 2
        for i in range(n_downsampling):  # add downsampling layers
            mult = 2 ** i
            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
                      norm_layer(ngf * mult * 2),
                      nn.ReLU(True)]

        mult = 2 ** n_downsampling
        for i in range(n_blocks):       # add ResNet blocks

            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]

        for i in range(n_downsampling):  # add upsampling layers
            mult = 2 ** (n_downsampling - i)
            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
                                         kernel_size=3, stride=2,
                                         padding=1, output_padding=1,
                                         bias=use_bias),
                      norm_layer(int(ngf * mult / 2)),
                      nn.ReLU(True)]
        model += [nn.ReflectionPad2d(3)]
        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
        model += [nn.Tanh()]

        self.model = nn.Sequential(*model)

    def forward(self, inp):
        """Standard forward"""
        return self.model(torch.cat(inp, dim=1))

class ResnetBlock(nn.Module):
    """Define a Resnet block"""

    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
        """Initialize the Resnet block

        A resnet block is a conv block with skip connections
        We construct a conv block with build_conv_block function,
        and implement skip connections in <forward> function.
        Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf
        """
        super(ResnetBlock, self).__init__()
        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)

    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
        """Construct a convolutional block.

        Parameters:
            dim (int)           -- the number of channels in the conv layer.
            padding_type (str)  -- the name of padding layer: reflect | replicate | zero
            norm_layer          -- normalization layer
            use_dropout (bool)  -- if use dropout layers.
            use_bias (bool)     -- if the conv layer uses bias or not

        Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))
        """
        conv_block = []
        p = 0
        if padding_type == 'reflect':
            conv_block += [nn.ReflectionPad2d(1)]
        elif padding_type == 'replicate':
            conv_block += [nn.ReplicationPad2d(1)]
        elif padding_type == 'zero':
            p = 1
        else:
            raise NotImplementedError('padding [%s] is not implemented' % padding_type)

        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]
        if use_dropout:
            conv_block += [nn.Dropout(0.5)]

        p = 0
        if padding_type == 'reflect':
            conv_block += [nn.ReflectionPad2d(1)]
        elif padding_type == 'replicate':
            conv_block += [nn.ReplicationPad2d(1)]
        elif padding_type == 'zero':
            p = 1
        else:
            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]

        return nn.Sequential(*conv_block)

    def forward(self, x):
        """Forward function (with skip connections)"""
        out = x + self.conv_block(x)  # add skip connections
        return out
        
class NLayerDiscriminator(nn.Module):
    """Defines a PatchGAN discriminator"""

    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
        """Construct a PatchGAN discriminator

        Parameters:
            input_nc (int)  -- the number of channels in input images
            ndf (int)       -- the number of filters in the last conv layer
            n_layers (int)  -- the number of conv layers in the discriminator
            norm_layer      -- normalization layer
        """
        super(NLayerDiscriminator, self).__init__()
        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
            use_bias = norm_layer.func == nn.InstanceNorm2d
        else:
            use_bias = norm_layer == nn.InstanceNorm2d
        kw = 4
        padw = 1
        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
        nf_mult = 1
        nf_mult_prev = 1
        for n in range(1, n_layers):  # gradually increase the number of filters
            nf_mult_prev = nf_mult
            nf_mult = min(2 ** n, 8)
            sequence += [
                        nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
                        norm_layer(ndf * nf_mult),
                        nn.LeakyReLU(0.2, True)
                        ]
        nf_mult_prev = nf_mult
        nf_mult = min(2 ** n_layers, 8)
        sequence += [
                    nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
                    norm_layer(ndf * nf_mult),
                    nn.LeakyReLU(0.2, True)
                    ]
        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]
        self.model = nn.Sequential(*sequence)

    def forward(self, input):
        """Standard forward."""
        return self.model(input)

class NLayerDiscriminatorS(nn.Module):
    """Defines a PatchGAN discriminator"""

    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
        """Construct a PatchGAN discriminator

        Parameters:
            input_nc (int)  -- the number of channels in input images
            ndf (int)       -- the number of filters in the last conv layer
            n_layers (int)  -- the number of conv layers in the discriminator
            norm_layer      -- normalization layer
        """
        super(NLayerDiscriminatorS, self).__init__()
        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
            use_bias = norm_layer.func == nn.InstanceNorm2d
        else:
            use_bias = norm_layer == nn.InstanceNorm2d
        kw = 4
        padw = 1
        sequence = [spectral_norm(nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw)), nn.LeakyReLU(0.2, True)]
        nf_mult = 1
        nf_mult_prev = 1
        for n in range(1, n_layers):  # gradually increase the number of filters
            nf_mult_prev = nf_mult
            nf_mult = min(2 ** n, 8)
            sequence += [
                        spectral_norm(nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias)),
                        norm_layer(ndf * nf_mult),
                        nn.LeakyReLU(0.2, True)
                        ]
        nf_mult_prev = nf_mult
        nf_mult = min(2 ** n_layers, 8)
        sequence += [
                    spectral_norm(nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias)),
                    norm_layer(ndf * nf_mult),
                    nn.LeakyReLU(0.2, True)
                    ]
        sequence += [spectral_norm(nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw))]
        self.model = nn.Sequential(*sequence)

    def forward(self, input):
        """Standard forward."""
        return self.model(input)    

######### MFNet ##########
class Local_Atten(nn.Module):
    def __init__(self, ngf=64):
        super(Local_Atten, self).__init__()
        self.self_atten = Self_Attn(ngf*4)
        self.attention = nn.Linear(ngf*4, 100)
        self.context_vec = nn.Linear(100, 1, bias=False)
        self.softmax  = nn.Softmax(dim=1) 
        
    def forward(self, style_features):
        B, C, H, W= style_features.shape
        h = self.self_atten(style_features)
        h = h.permute(0, 2, 3, 1).reshape(-1, C)
        h = torch.tanh(self.attention(h))                                   # (B*H*W, 100)
        h = self.context_vec(h)                                             # (B*H*W, 1)
        attention_map = self.softmax(h.view(B, H*W)).view(B, 1, H, W)       # (B, 1, H, W) 
        style_features = torch.sum(style_features*attention_map, dim=[2, 3])
        return style_features
    
class Layer_Atten(nn.Module):
    def __init__(self, ngf=64):
        super(Layer_Atten, self).__init__()
        self.ngf = ngf
        self.fc = nn.Linear(4096, 3)
        self.softmax  = nn.Softmax(dim=1) 
        
    def forward(self, style_features, style_features_1, style_features_2, style_features_3, B, K):
        
        style_features = torch.mean(style_features.view(B, K, self.ngf*4, 4, 4), dim=1)
        style_features = style_features.view(B, -1)
        weight = self.softmax(self.fc(style_features))
        
        style_features_1 = torch.mean(style_features_1.view(B, K, self.ngf*4), dim=1)
        style_features_2 = torch.mean(style_features_2.view(B, K, self.ngf*4), dim=1)
        style_features_3 = torch.mean(style_features_3.view(B, K, self.ngf*4), dim=1)
        
        style_features = (style_features_1*weight.narrow(1, 0, 1)+
                          style_features_2*weight.narrow(1, 1, 1)+
                          style_features_3*weight.narrow(1, 2, 1)).view(B, self.ngf*4, 1, 1)+torch.randn([B, self.ngf*4, 16, 16], device='cuda')*0.02
        return style_features

class MFNet_Style_Encoder(nn.Module):
    def __init__(self, input_nc, ngf=64):
        super(MFNet_Style_Encoder, self).__init__()
        # Conv Layers
        conv_model = [nn.ReflectionPad2d(3),
                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=False),
                 nn.BatchNorm2d(ngf),
                 nn.ReLU(True)]
        
        for i in range(2):  # add downsampling layers
            mult = 2 ** i
            conv_model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=False),
                      nn.BatchNorm2d(ngf * mult * 2),
                      nn.ReLU(True)
                      ]
        self.conv_model = nn.Sequential(*conv_model)
        
        # Multi-layer Attention Module
        self.local_atten_1 = Local_Atten(ngf=ngf)
        self.local_atten_2 = Local_Atten(ngf=ngf)
        self.local_atten_3 = Local_Atten(ngf=ngf)
        self.layer_atten = Layer_Atten(ngf=ngf)
        
        self.downsample_1 = nn.Sequential(nn.Conv2d(ngf * 4, ngf * 4, kernel_size=3, stride=2, padding=1, bias=False),
                          nn.BatchNorm2d(ngf * 4),
                          nn.ReLU(True)
                          )
        self.downsample_2 = nn.Sequential(nn.Conv2d(ngf * 4, ngf * 4, kernel_size=3, stride=2, padding=1, bias=False),
                          nn.BatchNorm2d(ngf * 4),
                          nn.ReLU(True)
                          )
        
    def forward(self, inp, B, K):
        # Results from conv layer
        style_features = self.conv_model(inp)
        ## local attention
        style_features_1 = self.local_atten_1(style_features)
        style_features = self.downsample_1(style_features)
        style_features_2 = self.local_atten_2(style_features)
        style_features = self.downsample_2(style_features)
        style_features_3 = self.local_atten_3(style_features)
        ## layer attention
        return self.layer_atten(style_features, style_features_1, style_features_2, style_features_3, B, K)

class MFNet_Content_Encoder(nn.Module):
    def __init__(self, input_nc, ngf=64):
        super(MFNet_Content_Encoder, self).__init__()
        # Conv Layers
        self.rf1 = nn.ReflectionPad2d(3)
        self.conv1 = nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=False)
        self.bn1 = nn.BatchNorm2d(ngf)
        self.relu1 = nn.ReLU(True)
        
        self.conv2 = nn.Conv2d(ngf, ngf * 2, kernel_size=3, stride=2, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(ngf * 2)
        self.relu2 = nn.ReLU(True)
        
        self.conv3 = nn.Conv2d(ngf * 2, ngf * 4, kernel_size=3, stride=2, padding=1, bias=False)
        self.bn3 = nn.BatchNorm2d(ngf * 4)
        self.relu3 = nn.ReLU(True)
        
    def forward(self, inp):
        x = self.rf1(inp)
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        skip1 = x   # channel = ngf * 2
        
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        skip2 = x   # channel = ngf * 4
        
        x = self.conv3(x)
        x = self.bn3(x)
        x = self.relu3(x)
    
        return x, skip1, skip2

class MFNet_Decoder(nn.Module):
    def __init__(self, use_dropout=False, n_blocks=6, ngf=64):
        super(MFNet_Decoder, self).__init__()
        res_blocks = []
        for i in range(n_blocks):       # add ResNet blocks
            res_blocks += [ResnetBlock(ngf*8, padding_type='reflect', norm_layer=nn.BatchNorm2d, use_dropout=use_dropout, use_bias=False)]
        self.res = nn.Sequential(*res_blocks)
        self.cont1 = nn.ConvTranspose2d(ngf*8, int(ngf*2),
                                         kernel_size=3, stride=2,
                                         padding=1, output_padding=1,
                                         bias=False)
        self.bn1 = nn.BatchNorm2d(int(ngf*2))
        self.relu1 = nn.ReLU(True)
        
        self.cont2 = nn.ConvTranspose2d(ngf*4, int(ngf),
                                         kernel_size=3, stride=2,
                                         padding=1, output_padding=1,
                                         bias=False)
        self.bn2 = nn.BatchNorm2d(int(ngf))
        self.relu2 = nn.ReLU(True)
        
        self.cont3 = nn.ConvTranspose2d(ngf*2, int(ngf),
                                         kernel_size=3, stride=2,
                                         padding=1, output_padding=1,
                                         bias=False)
        self.bn3 = nn.BatchNorm2d(int(ngf))
        self.relu3 = nn.ReLU(True)
        
        self.rp4 = nn.ReflectionPad2d(3)
        self.cont4 = nn.Conv2d(ngf, 1, kernel_size=7, padding=3, stride=2)
        self.tanh4 = nn.Tanh()
    
    def forward(self, inp, skip1, skip2):
        x = self.res(inp)
        x = self.cont1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.cont2(torch.cat((x, skip2), dim=1))
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.cont3(torch.cat((x, skip1), dim=1))
        x = self.bn3(x)
        x = self.relu3(x)
        x = self.cont4(x)
        x = self.tanh4(x)
        return x

class LCA_Skip_Connection(nn.Module):
	def __init__(self, in_dim1, in_dim2, tex_dim=32, str_dim=32):
		super(LCA_Skip_Connection, self).__init__()
		self.in_dim1 = in_dim1
		self.in_dim2 = in_dim2

		self.tex_dim = tex_dim
		self.str_dim = str_dim
		
		self.linear1 = nn.Sequential(
			nn.Linear(self.in_dim1**2, 256),
			nn.LeakyReLU(0.2, inplace=True),
			nn.Linear(256, tex_dim),
		)

		self.linear2 = nn.Sequential(
			nn.Linear(self.in_dim2**2, 256),
			nn.LeakyReLU(0.2, inplace=True),
			nn.Linear(256, tex_dim),
		)

		self.conv1 = nn.Sequential(
			nn.Conv2d(in_dim1, str_dim, kernel_size=1, padding=0, stride=2, bias=False),
			nn.BatchNorm2d(str_dim),
			nn.ReLU(True),
			nn.AvgPool2d(3, stride=2, padding=1)
		)

		self.conv2 = nn.Sequential(
			nn.Conv2d(in_dim2, str_dim, kernel_size=1, padding=0, bias=False),
			nn.BatchNorm2d(str_dim),
			nn.ReLU(True),
			nn.AvgPool2d(3, stride=2, padding=1)
		)

		self.conv_ind = nn.Sequential(
			nn.Conv2d(tex_dim+str_dim, 4, kernel_size=3, stride=2, padding=1, bias=False),
			nn.BatchNorm2d(4),
			nn.ReLU(True),
			nn.Conv2d(4, 1, kernel_size=3, padding=1, bias=False),
			nn.BatchNorm2d(1),
			nn.ReLU(True),
			nn.AvgPool2d(3, stride=2, padding=1)
		)

		self.linear_ind = nn.Sequential(
			nn.Linear(4*4, 1),
			nn.Sigmoid(),
		)

	def gaussian_blur(self, x, alpha):
		B, C, W, H = x.size()
		results = []
		
		if B > 1:
			alpha = alpha.squeeze()

		for batch in range(B):
			kernel_size = int(int(int(alpha[batch].item()*8)/2)*2)+1
		
			channels = C
			x_cord = torch.arange(kernel_size+0.)
			x_grid = x_cord.repeat(kernel_size).view(kernel_size, kernel_size)
			y_grid = x_grid.t()
			xy_grid = torch.stack([x_grid, y_grid], dim=-1)
			mean = (kernel_size - 1)//2
			diff = -torch.sum((xy_grid - mean)**2., dim=-1)
			gaussian_filter = torch.nn.Conv2d(in_channels=channels, out_channels=channels,
										kernel_size=kernel_size, groups=channels, bias=False)

			gaussian_filter.weight.requires_grad = False

			sigma = 16
			variance = sigma**2.
			gaussian_kernel = (1./(2.*math.pi*variance)) * torch.exp(diff /(2*variance))
			gaussian_kernel = gaussian_kernel / torch.sum(gaussian_kernel)
			gaussian_kernel = gaussian_kernel.view(1, 1, kernel_size, kernel_size)
			gaussian_kernel = gaussian_kernel.repeat(channels, 1, 1, 1)
			gaussian_kernel = gaussian_kernel.cuda()
			gaussian_filter.weight.data = gaussian_kernel
			output = gaussian_filter(torch.nn.functional.pad(x[batch].unsqueeze(0), (mean ,mean, mean, mean), "replicate"))

			results.append(output)
			
		return  torch.stack(results).squeeze(1)

	def gram_matrix(self, y):
		(b, ch, h, w) = y.size()
		features = y.view(b, ch, w * h)
		features_t = features.transpose(1, 2)
		gram = features.bmm(features_t) / (ch * h * w)
		return gram

	def forward(self, skip1, skip2):
		texture_feat_1 = self.linear1(self.gram_matrix(skip1).view(skip1.size(0),-1)).unsqueeze(2).unsqueeze(3).repeat(1, 1, 16, 16)
		structure_fear_1 = self.conv1(skip1)
		feat_1 = torch.cat([texture_feat_1, structure_fear_1], dim=1)

		texture_feat_2 = self.linear2(self.gram_matrix(skip2).view(skip2.size(0),-1)).unsqueeze(2).unsqueeze(3).repeat(1, 1, 16, 16)
		structure_fear_2 = self.conv2(skip2)
		feat_2 = torch.cat([texture_feat_2, structure_fear_2], dim=1)

		kernel1 = self.linear_ind(self.conv_ind(feat_1).view(skip1.size(0),-1))
		kernel2 = self.linear_ind(self.conv_ind(feat_2).view(skip2.size(0),-1))

		return self.gaussian_blur(skip1, kernel1), self.gaussian_blur(skip2, kernel2), kernel1, kernel2


class MFNet_Generator(nn.Module):
    def __init__(self, ngf=64, use_dropout=False, n_blocks=6):
        super(MFNet_Generator, self).__init__()
        self.style_encoder = MFNet_Style_Encoder(input_nc=1, ngf=ngf)
        self.content_encoder = MFNet_Content_Encoder(input_nc=1, ngf=ngf)
        self.decoder = MFNet_Decoder(use_dropout=use_dropout, n_blocks=n_blocks, ngf=ngf)
        self.LCA_skip_connection = LCA_Skip_Connection(64, 128)

    def forward(self, inp):
        cnt_img, sty_imgs = inp
        B_c, K_c, _, _ = cnt_img.shape
        B_s, K_s, _, _ = sty_imgs.shape

        # Get the content feature
        cnt_fea, skip1, skip2 = self.content_encoder(cnt_img)
        ## Language-complexity aware skip connection
        skip1, skip2, kernel1, kernel2 = self.LCA_skip_connection(skip1, skip2)

        # Get the style feature
        sty_fea = self.style_encoder(sty_imgs.view(-1, 1, 64, 64), B_s, K_s)
        
        # Get the generated image
        fea = torch.cat([cnt_fea, sty_fea], dim=1)
        gen_img = self.decoder(fea, skip1, skip2)

        # For loss calculation
        ## Get the encoded feature from generated image
        cnt_fea_fake, _, _ = self.content_encoder(gen_img)
        sty_fea_fake = self.style_encoder(gen_img, B_c, K_c)
        
        return gen_img, cnt_fea, cnt_fea_fake, sty_fea, sty_fea_fake, kernel1, kernel2


######### SA-GAN #########
class Self_Attn(nn.Module):
    """ Self attention Layer"""
    def __init__(self,in_dim):
        super(Self_Attn,self).__init__()
        self.chanel_in = in_dim
        self.query_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)
        self.key_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)
        self.value_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim , kernel_size= 1)
        self.gamma = nn.Parameter(torch.zeros(1))
        self.softmax  = nn.Softmax(dim=-1) 
        
    def forward(self,x):
        """
            inputs :
                x : input feature maps( B X C X W X H)
            returns :
                out : self attention value + input feature 
                attention: B X N X N (N is Width*Height)
        """
        m_batchsize,C,width ,height = x.size()
        proj_query  = self.query_conv(x).view(m_batchsize,-1,width*height).permute(0,2,1) # B X CX(N)
        proj_key =  self.key_conv(x).view(m_batchsize,-1,width*height) # B X C x (*W*H)
        energy =  torch.bmm(proj_query,proj_key) # transpose check
        attention = self.softmax(energy) # BX (N) X (N) 
        proj_value = self.value_conv(x).view(m_batchsize,-1,width*height) # B X C X N

        out = torch.bmm(proj_value,attention.permute(0,2,1) )
        out = out.view(m_batchsize,C,width,height)
        
        out = self.gamma*out + x
        return out
