from networks.auto_encoder import AutoEncoder
from torch.nn.functional import mse_loss
from torch.nn import Sequential, Linear, ReLU, Dropout, Sigmoid
from torch.distributions import Normal
from networks.modules import *
import torch


class AdversarialAutoEncoder(AutoEncoder):

    def __init__(self, *args, **kwargs):
        super(AdversarialAutoEncoder, self).__init__(*args, **kwargs)
        self.discriminator = Discriminator(self.latent_dim, self.dataParams)

    def forward(self, batch):
        # Encoder
        #  outputs, hidden (Batch, Timesteps aka. Size, Features / Latent Dim Size)
        z = self.encoder(batch)
        # Decoder
        # First repeat the data accordingly to the batch size
        z_repeatet = Repeater((batch.shape[0], self.dataParams['size'], -1))(z)
        x_hat = self.decoder(z_repeatet)
        return z, x_hat


class AdversarialAELightningOverrides:

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

    def forward(self, x):
        return self.network.forward(x)

    def training_step(self, batch, _, optimizer_i):
        if optimizer_i == 0:
            # ---------------------
            #  Train Discriminator
            # ---------------------
            # latent_fake, reconstruction
            latent_fake, _ = self.network.forward(batch)
            latent_real = self.normal.sample(latent_fake.shape)

            # Evaluate the input
            d_real_prediction = self.network.discriminator.forward(latent_real)
            d_fake_prediction = self.network.discriminator.forward(latent_fake)

            # Train the discriminator
            d_loss_real = mse_loss(d_real_prediction, torch.zeros(d_real_prediction.shape))
            d_loss_fake = mse_loss(d_fake_prediction, torch.ones(d_fake_prediction.shape))

            # Calculate the mean over both the real and the fake acc
            # ToDo: do i need to compute this seperate?
            d_loss = 0.5 * torch.add(d_loss_real, d_loss_fake)
            return {'loss': d_loss}

        elif optimizer_i == 1:
            # ---------------------
            #  Train AutoEncoder
            # ---------------------
            # z, x_hat
            _, batch_hat = self.forward(batch)
            loss = mse_loss(batch, batch_hat)
            return {'loss': loss}

        else:
            raise RuntimeError('This should not have happened, catch me if u can.')


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