"""
PyTorch implementation of a Latent Gaussian Model (LGM).

This model implements an encoder-decoder pair with a latent space that learns
a generative model of the input distribution using Gaussian distributions. The
model exposes the standard forward method plus a few extrans to allow quick
sampling of the learned prior images and the posterior over latent variables
given an input. The usur can also get the mean embedding (discarding the
parameters using the LGM.embed function.

This is the model most commonly associated with VAEs. However we have opted
for the nomenclature in Rezende et. al, 2014 since this decouples the loss
from the architecture, which techincally can be used with non-variational
losses (as with the Wasserstein Autoencoder).
"""


import torch.nn as nn
import torch
from .stochastic import DiagonalGaussian, HomoscedasticGaussian
from .initialization import weights_init


def get_latent(latent_type):
    if latent_type == 'diagonal':
        return DiagonalGaussian
    elif latent_type == 'homoscedastic':
        return HomoscedasticGaussian
    raise ValueError('Unrecognized latent layer {}'.format(latent_type))


class LGM(nn.Module):
    def __init__(self, latent_size, encoder, decoder, latent_type='diagonal'):
        super().__init__()
        self.encoder = encoder
        self.latent = get_latent(latent_type)(latent_size)
        self.decoder = decoder

        self.reset_parameter()

    def reset_parameter(self):
        self.apply(weights_init)

    @property
    def nlayers(self):
        return len(self.encoder)

    @property
    def latent_size(self):
        return self.latent.size

    def posterior(self, inputs):
        """
        Return the posterior distribution given the inputs
        """
        h = self.encoder(inputs)
        return self.latent(h)[1]

    def decode(self, z):
        return self.decoder(z)

    def embed(self, inputs):
        """Embed a batch of data points, x, into their z representations."""
        h = self.encoder(inputs)
        return self.latent(h)[0]

    def forward(self, inputs):
        """
        Takes a batch of samples, encodes them, and then decodes them again.
        Returns the parameters of the posterior to enable ELBO computation.
        """
        h = self.encoder(inputs)
        z, z_params = self.latent(h)  # z_params = (mu, logvar)
        return self.decoder(z), z, z_params

    def sample(self, inputs=None, n_samples=1):
        """
        Sample from the prior distribution or the conditional posterior
        learned by the model. If no input is given the output will have
        size (n_samples, latent_size) and if mu and logvar are given it
        will have size (batch_size, n_samples, latent size)
        """
        zsamples = self.sample_latent(inputs, n_samples)
        return self.decode(zsamples)

    def sample_latent(self, inputs=None, n_samples=1):
        if inputs is None:
            h = self.encoder[-1].bias.new_zeros(1, 2 * self.latent_size)
        else:
            h = self.encoder(inputs)
        return self.latent.sample(h, n_samples)


        
class pLGM(nn.Module):
    def __init__(self, latent_size, encoder, decoder):
        super().__init__()
        self.encoder = encoder
        self._latent_size = latent_size
        self.decoder = decoder

        self.reset_parameter()

    def reset_parameter(self):
        self.apply(weights_init)

    @property
    def nlayers(self):
        return len(self.encoder)

    @property
    def latent_size(self):
        return self._latent_size

    def posterior(self, inputs):
        """
        Return the posterior distribution given the inputs
        """
        h = self.encoder(inputs)
        return h
    
    def spread(self, latent, batch_size, h, w):
        spreaded = latent.view(batch_size, -1, 1, 1).expand(batch_size, -1, h, w)
        return spreaded
    
    def decode(self, z, img):
        batch_size, _, h, w = img.shape
        spreaded = self.spread(z, batch_size, h, w)
        img_plus = torch.cat([img, spreaded], dim = 1)
        return self.decoder(img_plus)

    def embed(self, inputs):
        """Embed a batch of data points, x, into their z representations."""
        h = self.encoder(inputs)
        return h

    def forward(self, inputs):
        """
        Takes a batch of samples, encodes them, and then decodes them again.
        Returns the parameters of the posterior to enable ELBO computation.
        """
        h = self.encoder(inputs)
        #z, z_params = self.latent(h)  # z_params = (mu, logvar)
        return self.decoder(h), h

epsilon = 1e-9
class sopLGM(nn.Module):
    def __init__(self, latent_size, so_n_repre, encoder, decoder):
        super().__init__()
        self.encoder = encoder
        self.latent_size = latent_size
        self.decoder = decoder
        self.so_n_repre = so_n_repre
        self.reset_parameter()

    def reset_parameter(self):
        self.apply(weights_init)

    def encode(self, inputs):
        z = self.encoder(inputs)
        z_linear = z[:,:self.latent_size]
        z_so_n = z[:,self.latent_size:]
        cs_list, matrix_list = self.so_n_repre.forward(z_so_n)
        return z, z_linear, z_so_n, cs_list, matrix_list

    def spread(self, latent, batch_size, h, w):
        spreaded = latent.view(batch_size, -1, 1, 1).expand(batch_size, -1, h, w)
        return spreaded
    
    def decode(self, z, img):
        batch_size, _, h, w = img.shape
        spreaded = self.spread(z, batch_size, h, w)
        img_plus = torch.cat([img, spreaded], dim = 1)
        return self.decoder(img_plus)

    def z_matrix_to_decoded(self, z_linear, matrix_list, img):
        matrix_flatten_list = [matrix.flatten(start_dim = 1) for matrix in matrix_list]
        z=torch.cat([z_linear] + matrix_flatten_list, dim = 1)
        decoded = self.decode(z, img)
        return decoded

    def forward(self, img_pair, img):
        #z = self.encoder(img_pair)
        #z_linear = z[:,:self.latent_size]
        #z_so_n = z[:,self.latent_size:]
        #cs_list, matrix_list = self.so_n_repre.forward(z_so_n)
        z, z_linear, z_so_n, cs_list, matrix_list = self.encode(img_pair)
        decoded = self.z_matrix_to_decoded(z_linear, matrix_list, img)
        return decoded, z_linear, cs_list, matrix_list

    @property
    def nlayers(self):
        return len(self.encoder)

