from torch.optim import Adam

from .modules import *
from torch.nn.functional import mse_loss


#######################
# Basic AE-Implementation
class VariationalAE(AbstractNeuralNetwork, ABC):

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

    def __init__(self, latent_dim=0, features=0, use_norm=True, train_on_predictions=False, **kwargs):
        assert latent_dim and features
        super(VariationalAE, self).__init__()
        self.features = features
        self.latent_dim = latent_dim
        self.encoder = Encoder(self.latent_dim, variational=True, use_norm=use_norm)
        self.decoder = Decoder(self.latent_dim, self.features, variational=True, use_norm=use_norm)

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

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

    def training_step(self, batch, _):
        x, y = batch
        mu, logvar, x_hat = self.forward(x)
        BCE = mse_loss(y, x_hat) if self.train_on_predictions else mse_loss(x, x_hat)

        # 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}

    def configure_optimizers(self):
        return [Adam(self.parameters(), lr=0.004)]


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