from torch.nn import Sequential, Linear, GRU, ReLU
from .modules import *
from torch.nn.functional import mse_loss


#######################
# Basic AE-Implementation
class BasicVAE(Module, ABC):

    @property
    def name(self):
        return self.__class__.__name__

    def __init__(self, dataParams, **kwargs):
        super(BasicVAE, self).__init__()
        self.dataParams = dataParams
        self.latent_dim = kwargs.get('latent_dim', 2)
        self.encoder = self._build_encoder()
        self.decoder = self._build_decoder(out_shape=self.dataParams['features'])
        self.mu, self.logvar = Linear(10, self.latent_dim), Linear(10, self.latent_dim)

    def _build_encoder(self):
        linear_stack = Sequential(
            Linear(6, 100, bias=True),
            ReLU(),
            Linear(100, 10, bias=True),
            ReLU()
        )
        encoder = Sequential(
            TimeDistributed(linear_stack),
            GRU(10, 10, batch_first=True),
            RNNOutputFilter(only_last=True),
        )
        return encoder

    def reparameterize(self, mu, logvar):
        # Lambda Layer, add gaussian noise
        std = torch.exp(0.5*logvar)
        eps = torch.randn_like(std)
        return mu + eps*std

    def _build_decoder(self, out_shape):
        decoder = Sequential(
            Linear(10, 100, bias=True),
            ReLU(),
            Linear(100, out_shape, bias=True),
            ReLU()
        )

        sequential_decoder = Sequential(
            GRU(self.latent_dim, 10, batch_first=True),
            RNNOutputFilter(),
            TimeDistributed(decoder)
        )
        return sequential_decoder

    def forward(self, batch):
        encoding = self.encoder(batch)
        mu_logvar = self.mu(encoding), self.logvar(encoding)
        z = self.reparameterize(*mu_logvar)
        repeat = Repeater((batch.shape[0], self.dataParams['size'], -1))
        x_hat = self.decoder(repeat(z))
        return (x_hat, *mu_logvar)


class VAELightningOverrides:

    def training_step(self, x, batch_nb):
        x_hat, logvar, mu = self.forward(x)
        BCE = mse_loss(x_hat, x, reduction='mean')

        # see Appendix B from VAE paper:
        # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
        # https://arxiv.org/abs/1312.6114
        # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
        KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        return {'loss': BCE + KLD}


if __name__ == '__main__':
    raise PermissionError('Get out of here - never run this module')