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, 20, 148, 91), 'torchvision.datasets.CelebA', 'CelebA', (), '', False, 'from torchvision.datasets import CelebA\n'), ((149, 18, 149, 88), 'torchvision.datasets.CelebA', 'CelebA', (), '', False, 'from torchvision.datasets import CelebA\n'), ((151, 19, 155, 45), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader\n'), ((156, 17, 160, 43), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader\n'), ((166, 16, 166, 103), 'pytorch_lightning.loggers.TensorBoardLogger', 'TensorBoardLogger', (), '', False, 'from pytorch_lightning.loggers import TensorBoardLogger\n'), ((167, 14, 167, 68), 'pytorch_lightning.Trainer', 'pl.Trainer', (), '', True, 'import pytorch_lightning as pl\n'), ((31, 24, 31, 44), 'torch.nn.Linear', 'nn.Linear', ({(31, 34, 31, 38): '4096', (31, 40, 31, 43): '200'}, {}), '(4096, 200)', True, 'import torch.nn as nn\n'), ((32, 28, 32, 48), 'torch.nn.Linear', 'nn.Linear', ({(32, 38, 32, 42): '4096', (32, 44, 32, 47): '200'}, {}), '(4096, 200)', True, 'import torch.nn as nn\n'), ((49, 29, 49, 49), 'torch.nn.Linear', 'nn.Linear', ({(49, 39, 49, 42): '200', (49, 44, 49, 48): '4096'}, {}), '(200, 4096)', True, 'import torch.nn as nn\n'), ((84, 14, 84, 37), 'torch.exp', 'torch.exp', ({(84, 24, 84, 36): '0.5 * logvar'}, {}), '(0.5 * logvar)', False, 'import torch\n'), ((85, 14, 85, 35), 'torch.randn_like', 'torch.randn_like', ({(85, 31, 85, 34): 'std'}, {}), '(std)', False, 'import torch\n'), ((100, 21, 100, 47), 'torch.nn.functional.mse_loss', 'F.mse_loss', ({(100, 32, 100, 41): 'recon_img', (100, 43, 100, 46): 'img'}, {}), '(recon_img, img)', True, 'import torch.nn.functional as F\n'), ((118, 21, 118, 47), 'torch.nn.functional.mse_loss', 'F.mse_loss', ({(118, 32, 118, 41): 'recon_img', (118, 43, 118, 46): 'img'}, {}), '(recon_img, img)', True, 'import torch.nn.functional as F\n'), ((134, 12, 134, 41), 'torch.randn', 'torch.randn', ({(134, 24, 134, 35): 'num_samples', (134, 37, 134, 40): '200'}, {}), '(num_samples, 200)', False, 'import torch\n'), ((17, 12, 17, 64), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((18, 12, 18, 30), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(18, 27, 18, 29): '32'}, {}), '(32)', True, 'import torch.nn as nn\n'), ((19, 12, 19, 26), 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((20, 12, 20, 65), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((21, 12, 21, 30), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(21, 27, 21, 29): '64'}, {}), '(64)', True, 'import torch.nn as nn\n'), ((22, 12, 22, 26), 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((23, 12, 23, 65), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((24, 12, 24, 30), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(24, 27, 24, 29): '64'}, {}), '(64)', True, 'import torch.nn as nn\n'), ((25, 12, 25, 26), 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((26, 12, 26, 65), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((27, 12, 27, 30), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(27, 27, 27, 29): '64'}, {}), '(64)', True, 'import torch.nn as nn\n'), ((28, 12, 28, 26), 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((37, 14, 37, 26), 'torch.nn.Flatten', 'nn.Flatten', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((52, 12, 52, 92), 'torch.nn.ConvTranspose2d', 'nn.ConvTranspose2d', (), '', True, 'import torch.nn as nn\n'), ((53, 12, 53, 30), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(53, 27, 53, 29): '64'}, {}), '(64)', True, 'import torch.nn as nn\n'), ((54, 12, 54, 26), 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((55, 12, 55, 92), 'torch.nn.ConvTranspose2d', 'nn.ConvTranspose2d', (), '', True, 'import torch.nn as nn\n'), ((56, 12, 56, 30), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(56, 27, 56, 29): '64'}, {}), '(64)', True, 'import torch.nn as nn\n'), ((57, 12, 57, 26), 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((58, 12, 58, 92), 'torch.nn.ConvTranspose2d', 'nn.ConvTranspose2d', (), '', True, 'import torch.nn as nn\n'), ((59, 12, 59, 30), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(59, 27, 59, 29): '32'}, {}), '(32)', True, 'import torch.nn as nn\n'), ((60, 12, 60, 26), 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((61, 12, 61, 91), 'torch.nn.ConvTranspose2d', 'nn.ConvTranspose2d', (), '', True, 'import torch.nn as nn\n'), ((62, 12, 62, 29), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(62, 27, 62, 28): '3'}, {}), '(3)', True, 'import torch.nn as nn\n'), ((63, 12, 63, 24), 'torch.nn.Sigmoid', 'nn.Sigmoid', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((142, 8, 142, 41), 'torchvision.transforms.RandomHorizontalFlip', 'transforms.RandomHorizontalFlip', ({}, {}), '()', False, 'from torchvision import transforms\n'), ((143, 8, 143, 34), 'torchvision.transforms.CenterCrop', 'transforms.CenterCrop', ({(143, 30, 143, 33): '148'}, {}), '(148)', False, 'from torchvision import transforms\n'), ((144, 8, 144, 30), 'torchvision.transforms.Resize', 'transforms.Resize', ({(144, 26, 144, 29): '128'}, {}), '(128)', False, 'from torchvision import transforms\n'), ((145, 8, 145, 29), 'torchvision.transforms.ToTensor', 'transforms.ToTensor', ({}, {}), '()', False, 'from torchvision import transforms\n')]