import numpy as np
import torch
import functools
import torch.nn as nn
from torch.autograd import Variable


###############################################################################
# Functions
###############################################################################
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm2d') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)


def get_norm_layer(norm_type='instance'):
    if norm_type == 'batch':
        norm_layer = functools.partial(nn.BatchNorm2d, affine=True)
    elif norm_type == 'instance':
        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False)
    else:
        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
    return norm_layer


def define_G(input_nc, output_nc, n_phases, gpu_ids=[], res_blocks=1):
    netG = GlobalGenerator(input_nc, output_nc, n_phases, res_blocks=res_blocks)
    print(netG)
    if len(gpu_ids) > 0:
        assert (torch.cuda.is_available())
        netG.cuda(gpu_ids[0])
    netG.apply(weights_init)
    return netG


def define_GlueNet(input_nc, label_set, gpu_ids=[]):
    netGlue = GlueGenerator(input_nc, label_set)
    print(netGlue)
    if len(gpu_ids) > 0:
        assert (torch.cuda.is_available())
        netGlue.cuda(gpu_ids[0])
    netGlue.apply(weights_init)
    return netGlue


def define_D(input_nc, ndf, n_layers_D, norm='instance', use_sigmoid=False, num_D=2, getIntermFeat=False, gpu_ids=[]):
    norm_layer = get_norm_layer(norm_type=norm)
    netD = MultiscaleDiscriminator(input_nc, ndf, n_layers_D, norm_layer, use_sigmoid, num_D, getIntermFeat)
    print(netD)
    if len(gpu_ids) > 0:
        assert (torch.cuda.is_available())
        netD.cuda(gpu_ids[0])
    netD.apply(weights_init)
    return netD


def print_network(net):
    pass


################################################################################
# Losses
################################################################################
class GANLoss(nn.Module):
    def __init__(self, use_lsgan=True, target_real_label=1.0, target_fake_label=0.0,
                 tensor=torch.FloatTensor):
        super(GANLoss, self).__init__()
        self.real_label = target_real_label
        self.fake_label = target_fake_label
        self.real_label_var = None
        self.fake_label_var = None
        self.Tensor = tensor
        if use_lsgan:
            self.loss = nn.MSELoss()
        else:
            self.loss = nn.BCELoss()

    def get_target_tensor(self, input, target_is_real):
        target_tensor = None
        if target_is_real:
            create_label = ((self.real_label_var is None) or
                            (self.real_label_var.numel() != input.numel()))
            if create_label:
                real_tensor = self.Tensor(input.size()).fill_(self.real_label)
                self.real_label_var = Variable(real_tensor, requires_grad=False)
            target_tensor = self.real_label_var
        else:
            create_label = ((self.fake_label_var is None) or
                            (self.fake_label_var.numel() != input.numel()))
            if create_label:
                fake_tensor = self.Tensor(input.size()).fill_(self.fake_label)
                self.fake_label_var = Variable(fake_tensor, requires_grad=False)
            target_tensor = self.fake_label_var
        return target_tensor

    def __call__(self, input, target_is_real):
        if isinstance(input[0], list):
            loss = 0
            for input_i in input:
                pred = input_i[-1]
                target_tensor = self.get_target_tensor(pred, target_is_real)
                loss += self.loss(pred, target_tensor)
            return loss
        else:
            target_tensor = self.get_target_tensor(input[-1], target_is_real)
            return self.loss(input[-1], target_tensor)


class VGGLoss(nn.Module):
    def __init__(self, gpu_ids):
        super(VGGLoss, self).__init__()
        self.vgg = Vgg19().cuda()
        self.criterion = nn.L1Loss()
        self.weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]

    def forward(self, x, y):
        x_vgg, y_vgg = self.vgg(x), self.vgg(y)
        loss = 0
        for i in range(len(x_vgg)):
            loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())
        return loss


class LabelLoss():
    def __init__(self):
        pass

    def __call__(self, *args, **kwargs):
        pass


################################################################################
# Generators
################################################################################

# inside of here, network takes a 35 layer one hot, transforms it into a higher res 35 layer big daddy, then uses a separate network to transform it into a 1 layer meaty bad boy
# what
class GlobalGenerator(nn.Module):
    def __init__(self, input_nc, output_nc, n_phases, ngf=64, norm_layer=nn.BatchNorm2d, padding_type='reflect',
                 res_blocks=1):
        assert (n_phases > 0)
        super(GlobalGenerator, self).__init__()
        self.n_phases = n_phases
        self.models = nn.ModuleList()
        self.final_layers = []

        def get_model_block(ngf, mult, res_blocks, norm_layer=nn.BatchNorm2d, activation=nn.ReLU(True),
                            padding_type='reflect', use_scaling=True):
            # down
            if use_scaling:
                model = [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=4, stride=2, padding=2),
                         norm_layer(ngf * mult * 2), activation]
            else:
                model = []

            # res
            for i in range(res_blocks):
                model += [ResnetBlock(ngf * 2, padding_type=padding_type, activation=activation, norm_layer=norm_layer)]

            # up
            if use_scaling:
                model += [nn.ConvTranspose2d(ngf * mult * 2, ngf * mult, kernel_size=4, stride=2, padding=2),
                          norm_layer(ngf * mult), activation]

            return model

        activation = nn.LeakyReLU(inplace=True)

        ## add first layer 1 -> 64, which is kinda like a feature extractor?? Is first layer
        model = [nn.Conv2d(input_nc, ngf, kernel_size=3, padding=1), norm_layer(ngf), activation]
        out_layer = nn.Sequential(nn.Conv2d(ngf, output_nc, kernel_size=3, padding=1), nn.Tanh())
        self.models.append(nn.ModuleList(nn.ModuleList([nn.Sequential(*model), out_layer])))

        ## add layers for phases
        for n in range(n_phases - 1):
            mult = 1
            # mult = 2**n
            model = get_model_block(ngf, mult, res_blocks, norm_layer, activation, padding_type)
            # final 64 -> 1 layer
            self.models.append(nn.ModuleList(nn.ModuleList(
                [nn.Sequential(*model),
                 nn.Sequential(nn.Conv2d(ngf, output_nc, kernel_size=3, padding=1), nn.Tanh())])))

    def forward(self, input_img, end_phase, blend_prev):
        """

        :param input_img: Is it one layer or 35 layer?
        :param end_phase:
        :param blend_prev:
        :return:
        """
        assert end_phase <= self.n_phases, 'phase is: {}, max: {}'.format(end_phase, self.n_phases)
        # 64 -> 64, until time comes for out, then 64 -> 64 -> 1 with blend of previous ->1
        # start on phase 1
        # want to predict the x2 detailed image
        # pass it through:
        # input convs models [0][0]
        # model resnet blocks models [1][0]
        # here choice - if phase = 1, thru models [1][1]
        # if if phase = 2, go to x1, through models [2][0] then out models[2][1]
        # so start at phase 1
        # at the start, train only the previous layer, cause blend is 1
        # print('out at phase: ', n)
        # start at end_phase 1
        # train model[0] until you start blending

        # 35 -> 64
        output = self.models[0][0](input_img)

        for n in range(1, end_phase):
            model = self.models[n]
            output = model[0](output)

        # the output blending magick
        if self.n_phases == end_phase:
            end_phase = end_phase - 1
            # final_model = self.models[end_phase - 1]
            # output = final_model[1](output)

        model = self.models[end_phase]
        prev_model = self.models[end_phase - 1]
        output = blend_prev * prev_model[1](output) + (1 - blend_prev) * model[1](model[0](output))

        return output


class GlueGenerator(nn.module):
    def __init__(self, input_nc, label_set_global):
        self.label_set_global = label_set_global  # a list(?) of allowed label values
        self.num_labels = len(label_set_global)

        lin_in = 1024
        # encode the input label set so that the bangaroos are conditioned on it
        # the encoder input technically varies in length. It's possible to pad it with -1s or something if it's empty
        # but maybe there's a better way
        self.encoder = None
        self.model = nn.Conv2d(input_nc + 1, 1, kernel_size=3, stride=1, padding=1)    # a convnet or something
        self.label_predictor = nn.Linear(lin_in, self.num_labels)

    # receive 35 chan input
    def forward(self, input_img, label_set_local):

        # use something like StackedGAN to make an embedding of label_set_local to attach to the convnet input
        # use a conv net to look at a point
        # decide which class point is gonna be
        # by predicting a one hot, indicating which element of label set to choose from
        # do for all pixels
        result = None
        return result


class ResnetBlock(nn.Module):
    def __init__(self, dim, padding_type, norm_layer, activation=nn.ReLU(True), use_dropout=False):
        super(ResnetBlock, self).__init__()
        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, activation, use_dropout)

    def build_conv_block(self, dim, padding_type, norm_layer, activation, use_dropout):
        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),
                       norm_layer(dim),
                       activation]
        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),
                       norm_layer(dim)]

        return nn.Sequential(*conv_block)

    def forward(self, x):
        out = x + self.conv_block(x)
        return out


################################################################################
# Discriminators
################################################################################
class MultiscaleDiscriminator(nn.Module):
    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d,
                 use_sigmoid=False, num_D=3, getIntermFeat=False):
        super(MultiscaleDiscriminator, self).__init__()
        self.num_D = num_D
        self.n_layers = n_layers
        self.getIntermFeat = getIntermFeat

        for i in range(num_D):
            netD = NLayerDiscriminator(input_nc, ndf, n_layers, norm_layer, use_sigmoid, getIntermFeat)
            if getIntermFeat:
                for j in range(n_layers + 2):
                    setattr(self, 'scale' + str(i) + '_layer' + str(j), getattr(netD, 'model' + str(j)))
            else:
                setattr(self, 'layer' + str(i), netD.model)

        self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)

    def singleD_forward(self, model, input):
        if self.getIntermFeat:
            result = [input]
            for i in range(len(model)):
                result.append(model[i](result[-1]))
            return result[1:]
        else:
            return [model(input)]

    def forward(self, input):
        num_D = self.num_D
        result = []
        input_downsampled = input
        for i in range(num_D):
            if self.getIntermFeat:
                model = [getattr(self, 'scale' + str(num_D - 1 - i) + '_layer' + str(j)) for j in
                         range(self.n_layers + 2)]
            else:
                model = getattr(self, 'layer' + str(num_D - 1 - i))
            result.append(self.singleD_forward(model, input_downsampled))
            if i != (num_D - 1):
                input_downsampled = self.downsample(input_downsampled)
        return result


class NLayerDiscriminator(nn.Module):
    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_sigmoid=False, getIntermFeat=False):
        super(NLayerDiscriminator, self).__init__()
        self.getIntermFeat = getIntermFeat
        self.n_layers = n_layers

        kw = 4
        padw = int(np.ceil((kw - 1.0) / 2))
        sequence = [[nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]]

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

        nf_prev = nf
        nf = min(nf * 2, 512)
        sequence += [[
            nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=1, padding=padw),
            norm_layer(nf),
            nn.LeakyReLU(0.2, True)
        ]]

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

        if use_sigmoid:
            sequence += [[nn.Sigmoid()]]

        if getIntermFeat:
            for n in range(len(sequence)):
                setattr(self, 'model' + str(n), nn.Sequential(*sequence[n]))
        else:
            sequence_stream = []
            for n in range(len(sequence)):
                sequence_stream += sequence[n]
            self.model = nn.Sequential(*sequence_stream)

    def forward(self, input):
        if self.getIntermFeat:
            res = [input]
            for n in range(self.n_layers + 2):
                model = getattr(self, 'model' + str(n))
                res.append(model(res[-1]))
            return res[1:]
        else:
            return self.model(input)


################################################################################
# Zoo
################################################################################
from torchvision import models


class Vgg19(torch.nn.Module):
    def __init__(self, requires_grad=False):
        super(Vgg19, self).__init__()
        vgg_pretrained_features = models.vgg19(pretrained=True).features
        self.slice1 = torch.nn.Sequential()
        self.slice2 = torch.nn.Sequential()
        self.slice3 = torch.nn.Sequential()
        self.slice4 = torch.nn.Sequential()
        self.slice5 = torch.nn.Sequential()
        for x in range(2):
            self.slice1.add_module(str(x), vgg_pretrained_features[x])
        for x in range(2, 7):
            self.slice2.add_module(str(x), vgg_pretrained_features[x])
        for x in range(7, 12):
            self.slice3.add_module(str(x), vgg_pretrained_features[x])
        for x in range(12, 21):
            self.slice4.add_module(str(x), vgg_pretrained_features[x])
        for x in range(21, 30):
            self.slice5.add_module(str(x), vgg_pretrained_features[x])
        if not requires_grad:
            for param in self.parameters():
                param.requires_grad = False

    def forward(self, X):
        h_relu1 = self.slice1(X)
        h_relu2 = self.slice2(h_relu1)
        h_relu3 = self.slice3(h_relu2)
        h_relu4 = self.slice4(h_relu3)
        h_relu5 = self.slice5(h_relu4)
        out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
        return out
