from torch.optim import Adam

from networks.auto_encoder import AutoEncoder
from torch.nn.functional import mse_loss
from networks.modules import *
import torch


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


class AdversarialAutoEncoder(AutoEncoder):

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

    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], batch.shape[1], -1))(z)
        x_hat = self.decoder(z_repeatet)
        return z, x_hat


class AdversarialAELightningOverrides(LightningModuleOverrides):
    
    def __init__(self):
        super(AdversarialAELightningOverrides, self).__init__()

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

            # 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, device=device))
            d_loss_fake = mse_loss(d_fake_prediction, torch.ones(d_fake_prediction.shape, device=device))

            # 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) * 0.001
            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.')


    # This is Fucked up, why do i need to put an additional empty list here?
    def configure_optimizers(self):
        return [Adam(self.network.discriminator.parameters(), lr=0.02),
                Adam([*self.network.encoder.parameters(), *self.network.decoder.parameters()], lr=0.02), ],\
               []


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