import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from pytorch_lightning.loggers import TensorBoardLogger
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import CelebA


class Encoder(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv_layers = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(),
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(),
        )

        self.mu_layer = nn.Linear(4096, 200)
        self.logvar_layer = nn.Linear(4096, 200)

    def forward(self, imgs):
        out = self.conv_layers(imgs)

        out = nn.Flatten()(out)

        mu = self.mu_layer(out)
        logvar = self.logvar_layer(out)

        return mu, logvar


class Decoder(nn.Module):
    def __init__(self):
        super().__init__()

        self.decoder_input = nn.Linear(200, 4096)

        self.deconv_layers = nn.Sequential(
            nn.ConvTranspose2d(64, 64, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(),
            nn.ConvTranspose2d(64, 64, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(),
            nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(),
            nn.ConvTranspose2d(32, 3, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(3),
            nn.Sigmoid(),
        )

    def forward(self, z):
        out = self.decoder_input(z)
        out = out.view(-1, 64, 8, 8)
        recon_img = self.deconv_layers(out)
        return recon_img


class VanillaVAE(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()

    def forward(self, img):
        mu, logvar = self.encoder(img)
        return mu

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return eps * std + mu

    def configure_optimizers(self):
        optimizer = optim.Adam(self.parameters(), lr=0.005)
        return optimizer

    def training_step(self, train_batch, batch_idx):
        img, labels = train_batch

        mu, logvar = self.encoder(img)
        z = self.reparameterize(mu, logvar)
        recon_img = self.decoder(z)

        recon_loss_factor = 10000
        recon_loss = F.mse_loss(recon_img, img)
        kld_loss = torch.mean(-0.5 * torch.sum(1 + logvar - mu**2 - logvar.exp(), dim=1))
        loss = recon_loss_factor * recon_loss + kld_loss

        self.log('train/loss', loss)
        self.log('train/recon_loss', recon_loss)
        self.log('train/kl_loss', kld_loss)

        return loss

    def validation_step(self, val_batch, batch_idx):
        img, labels = val_batch

        mu, logvar = self.encoder(img)
        z = self.reparameterize(mu, logvar)
        recon_img = self.decoder(z)

        recon_loss_factor = 10000
        recon_loss = F.mse_loss(recon_img, img)
        kld_loss = torch.mean(-0.5 * torch.sum(1 + logvar - mu**2 - logvar.exp(), dim=1))
        loss = recon_loss_factor * recon_loss + kld_loss

        self.log('val/loss', loss)
        self.log('val/recon_loss', recon_loss)
        self.log('val/kl_loss', kld_loss)

        return loss

    def reconstruct(self, img):
        mu, _ = self.encoder(img)
        recon_img = self.decoder(mu)
        return recon_img

    def sample(self, num_samples=64):
        z = torch.randn(num_samples, 200)
        samples = self.decoder(z)
        return samples


if __name__ == '__main__':
    # data
    transform = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        transforms.CenterCrop(148),
        transforms.Resize(128),
        transforms.ToTensor()
    ])

    train_dataset = CelebA(root='data', split='train', transform=transform, download=False)
    val_dataset = CelebA(root='data', split='test', transform=transform, download=False)

    train_loader = DataLoader(train_dataset,
                              batch_size=32,
                              num_workers=8,
                              shuffle=True,
                              drop_last=True)
    val_loader = DataLoader(val_dataset,
                            batch_size=32,
                            num_workers=8,
                            shuffle=False,
                            drop_last=True)

    # model
    model = VanillaVAE()

    # training
    tb_logger = TensorBoardLogger('lightning_logs', name='vanilla_vae_celeba', default_hp_metric=False)
    trainer = pl.Trainer(gpus=[0], max_epochs=200, logger=tb_logger)
    trainer.fit(model, train_loader, val_loader)

[(148, 'torchvision.datasets.CelebA', 'CelebA', ([], {'root': '"""data"""', 'split': '"""train"""', 'transform': 'transform', 'download': '(False)'}), False, 'from torchvision.datasets import CelebA\n'), (149, 'torchvision.datasets.CelebA', 'CelebA', ([], {'root': '"""data"""', 'split': '"""test"""', 'transform': 'transform', 'download': '(False)'}), False, 'from torchvision.datasets import CelebA\n'), (151, 'torch.utils.data.DataLoader', 'DataLoader', (['train_dataset'], {'batch_size': '(32)', 'num_workers': '(8)', 'shuffle': '(True)', 'drop_last': '(True)'}), False, 'from torch.utils.data import DataLoader\n'), (156, 'torch.utils.data.DataLoader', 'DataLoader', (['val_dataset'], {'batch_size': '(32)', 'num_workers': '(8)', 'shuffle': '(False)', 'drop_last': '(True)'}), False, 'from torch.utils.data import DataLoader\n'), (166, 'pytorch_lightning.loggers.TensorBoardLogger', 'TensorBoardLogger', (['"""lightning_logs"""'], {'name': '"""vanilla_vae_celeba"""', 'default_hp_metric': '(False)'}), False, 'from pytorch_lightning.loggers import TensorBoardLogger\n'), (167, 'pytorch_lightning.Trainer', 'pl.Trainer', ([], {'gpus': '[0]', 'max_epochs': '(200)', 'logger': 'tb_logger'}), True, 'import pytorch_lightning as pl\n'), (31, 'torch.nn.Linear', 'nn.Linear', (['(4096)', '(200)'], {}), True, 'import torch.nn as nn\n'), (32, 'torch.nn.Linear', 'nn.Linear', (['(4096)', '(200)'], {}), True, 'import torch.nn as nn\n'), (49, 'torch.nn.Linear', 'nn.Linear', (['(200)', '(4096)'], {}), True, 'import torch.nn as nn\n'), (84, 'torch.exp', 'torch.exp', (['(0.5 * logvar)'], {}), False, 'import torch\n'), (85, 'torch.randn_like', 'torch.randn_like', (['std'], {}), False, 'import torch\n'), (100, 'torch.nn.functional.mse_loss', 'F.mse_loss', (['recon_img', 'img'], {}), True, 'import torch.nn.functional as F\n'), (118, 'torch.nn.functional.mse_loss', 'F.mse_loss', (['recon_img', 'img'], {}), True, 'import torch.nn.functional as F\n'), (134, 'torch.randn', 'torch.randn', (['num_samples', '(200)'], {}), False, 'import torch\n'), (17, 'torch.nn.Conv2d', 'nn.Conv2d', (['(3)', '(32)'], {'kernel_size': '(3)', 'stride': '(2)', 'padding': '(1)'}), True, 'import torch.nn as nn\n'), (18, 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(32)'], {}), True, 'import torch.nn as nn\n'), (19, 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ([], {}), True, 'import torch.nn as nn\n'), (20, 'torch.nn.Conv2d', 'nn.Conv2d', (['(32)', '(64)'], {'kernel_size': '(3)', 'stride': '(2)', 'padding': '(1)'}), True, 'import torch.nn as nn\n'), (21, 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(64)'], {}), True, 'import torch.nn as nn\n'), (22, 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ([], {}), True, 'import torch.nn as nn\n'), (23, 'torch.nn.Conv2d', 'nn.Conv2d', (['(64)', '(64)'], {'kernel_size': '(3)', 'stride': '(2)', 'padding': '(1)'}), True, 'import torch.nn as nn\n'), (24, 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(64)'], {}), True, 'import torch.nn as nn\n'), (25, 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ([], {}), True, 'import torch.nn as nn\n'), (26, 'torch.nn.Conv2d', 'nn.Conv2d', (['(64)', '(64)'], {'kernel_size': '(3)', 'stride': '(2)', 'padding': '(1)'}), True, 'import torch.nn as nn\n'), (27, 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(64)'], {}), True, 'import torch.nn as nn\n'), (28, 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ([], {}), True, 'import torch.nn as nn\n'), (37, 'torch.nn.Flatten', 'nn.Flatten', ([], {}), True, 'import torch.nn as nn\n'), (52, 'torch.nn.ConvTranspose2d', 'nn.ConvTranspose2d', (['(64)', '(64)'], {'kernel_size': '(3)', 'stride': '(2)', 'padding': '(1)', 'output_padding': '(1)'}), True, 'import torch.nn as nn\n'), (53, 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(64)'], {}), True, 'import torch.nn as nn\n'), (54, 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ([], {}), True, 'import torch.nn as nn\n'), (55, 'torch.nn.ConvTranspose2d', 'nn.ConvTranspose2d', (['(64)', '(64)'], {'kernel_size': '(3)', 'stride': '(2)', 'padding': '(1)', 'output_padding': '(1)'}), True, 'import torch.nn as nn\n'), (56, 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(64)'], {}), True, 'import torch.nn as nn\n'), (57, 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ([], {}), True, 'import torch.nn as nn\n'), (58, 'torch.nn.ConvTranspose2d', 'nn.ConvTranspose2d', (['(64)', '(32)'], {'kernel_size': '(3)', 'stride': '(2)', 'padding': '(1)', 'output_padding': '(1)'}), True, 'import torch.nn as nn\n'), (59, 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(32)'], {}), True, 'import torch.nn as nn\n'), (60, 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ([], {}), True, 'import torch.nn as nn\n'), (61, 'torch.nn.ConvTranspose2d', 'nn.ConvTranspose2d', (['(32)', '(3)'], {'kernel_size': '(3)', 'stride': '(2)', 'padding': '(1)', 'output_padding': '(1)'}), True, 'import torch.nn as nn\n'), (62, 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(3)'], {}), True, 'import torch.nn as nn\n'), (63, 'torch.nn.Sigmoid', 'nn.Sigmoid', ([], {}), True, 'import torch.nn as nn\n'), (142, 'torchvision.transforms.RandomHorizontalFlip', 'transforms.RandomHorizontalFlip', ([], {}), False, 'from torchvision import transforms\n'), (143, 'torchvision.transforms.CenterCrop', 'transforms.CenterCrop', (['(148)'], {}), False, 'from torchvision import transforms\n'), (144, 'torchvision.transforms.Resize', 'transforms.Resize', (['(128)'], {}), False, 'from torchvision import transforms\n'), (145, 'torchvision.transforms.ToTensor', 'transforms.ToTensor', ([], {}), False, 'from torchvision import transforms\n')]