import torch

import torch.nn as nn

import pytorch_pretrained_bert as Bert

# borrowed boilerplate vae code from https://github.com/AntixK/PyTorch-VAE

class VAE(Bert.modeling.BertPreTrainedModel):

    def __init__(self, config):

        super(VAE, self).__init__(config)

        self.unsuplist = config.unsupSize

        vaelatentdim = config.vaelatentdim

        vaeinchannels = config.vaeinchannels

        modules = []

        vaehidden = [config.poolingSize]

        self.linearFC = nn.Linear(config.hidden_size, config.poolingSize)

        self.activ =  nn.ReLU()

        # Build Encoder

        self.fc_mu = nn.Linear(vaehidden[-1], vaelatentdim)

        self.fc_var = nn.Linear(vaehidden[-1], vaelatentdim)

        # Build Decoder

        modules = []

        self.decoder1 = nn.Linear(vaelatentdim, vaehidden[-1])

        self.decoder2 = nn.Linear(vaehidden[-1],int( vaehidden[-1]))

        self.logSoftmax = nn.LogSoftmax(dim=1)

        self.linearOut =  nn.ModuleList([nn.Linear   (int( vaehidden[-1]), el[0]) for el in self.unsuplist])

        self.BetaD = config.BetaD

        self.apply(self.init_bert_weights)

    def encode(self, input: torch.Tensor) :

        """

        Encodes the input by passing through the encoder network

        and returns the latent codes.

        :param input: (Tensor) Input tensor to encoder [N x C x H x W]

        :return: (Tensor) List of latent codes

        """

        # result = self.activ (self.linearFC(input))

        mu = self.fc_mu(input)

        log_var = self.fc_var(input)

        return [mu, log_var]

    def decode(self, z: torch.Tensor) -> torch.Tensor:

        """

        Maps the given latent codes

        onto the image space.

        :param z: (Tensor) [B x D]

        :return: (Tensor) [B x C x H x W]

        """

        result = self.activ(self.decoder1(z))

        result = self.activ(self.decoder2(result))

        outs = []

        for outputiter , linoutnetwork in enumerate(self.linearOut):

            resout = self.logSoftmax(linoutnetwork(result))

            outs.append(resout)

        outs = torch.cat((outs), dim=1)

        return outs

    def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:

        """

        Reparameterization trick to sample from N(mu, var) from

        N(0,1).

        :param mu: (Tensor) Mean of the latent Gaussian [B x D]

        :param logvar: (Tensor) Standard deviation of the latent Gaussian [B x D]

        :return: (Tensor) [B x D]

        """

        std = torch.exp(0.5 * logvar)

        eps = torch.randn_like(std)

        return eps * std + mu

    def forward(self, input: torch.Tensor, label: torch.Tensor):

        if self.BetaD==False:

            mu, log_var = self.encode(input)

            z = self.reparameterize(mu, log_var)

            return  [self.decode(z), label, mu, log_var]

        else:

            mu, log_var = self.encode(input)

            z = self.reparameterize(mu, log_var)

            return  [self.decode(z), label, mu, log_var]

    def loss_function(self,dictout) -> dict:

        """

        Computes the VAE loss function.

        KL(N(\mu, \sigma), N(0, 1)) = \log \frac{1}{\sigma} + \frac{\sigma^2 + \mu^2}{2} - \frac{1}{2}

        :param args:

        :param kwargs:

        :return:

        """

        recons = dictout[0].transpose(1,0)

        input = dictout[1].transpose(1,0)

        mu = dictout[2]

        log_var = dictout[3]

        if self.BetaD==False:

            kld_weight = self.config.klpar # Account for the minibatch samples from the dataset

            reconsloss = 0

            startindx = 0

            outs = []

            labs = []

            for outputiter , output in enumerate(self.unsuplist):

                elementssize = output[0]

                chunkrecons = recons[startindx:startindx+elementssize].transpose(1,0)

                labels= input[outputiter]

                lossF  = nn.NLLLoss(reduction='none', ignore_index=-1)

                temploss = lossF(chunkrecons,labels).sum()

                reconsloss =reconsloss+ temploss

                outs.append(chunkrecons)

                labs.append(labels)

                startindx = startindx+elementssize

            kld_loss = torch.sum(-0.5 * torch.sum(1 + log_var - mu ** 2 - log_var.exp(), dim = 1), dim = 0)

            loss = (reconsloss + kld_weight * kld_loss)/len(dictout[0])

            if self.config.klpar<1:

                self.config.klpar = self.config.klpar + 1e-5

            return {'loss': loss, 'Reconstruction_Loss':reconsloss, 'KLD':-kld_loss, 'outs':outs, 'labs':labs}

        else:

            return 0

    def sample(self,

               num_samples:int,

               current_device: int, **kwargs) -> torch.Tensor:

        """

        Samples from the latent space and return the corresponding

        image space map.

        :param num_samples: (Int) Number of samples

        :param current_device: (Int) Device to run the model

        :return: (Tensor)

        """

        z = torch.randn(num_samples,

                        self.vaelatentdim)

        z = z.to(current_device)

        samples = self.decode(z)

        return samples

    def generate(self, x: torch.Tensor, **kwargs) -> torch.Tensor:

        """

        Given an input image x, returns the reconstructed image

        :param x: (Tensor) [B x C x H x W]

        :return: (Tensor) [B x C x H x W]

        """

        return self.forward(x)[0]