import argparse
import os
import numpy as np
import pandas as pd
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation as animation
from IPython.display import HTML
manualSeed=999
print(manualSeed)
random.seed(manualSeed)
torch.manual_seed(manualSeed)
# Root directory for dataset
dataroot = "data/celeba"
# Number of workers for dataloader
workers = 2
# Batch size during training
batch_size = 128
# Spatial size of training images. All images will be resized to this
#   size using a transformer.
image_size = 64
# Number of channels in the training images. For color images this is 3
nc = 3
# Size of z latent vector (i.e. size of generator input)
nz = 100
# Size of feature maps in generator
ngf = 64
# Size of feature maps in discriminator
ndf = 64
# Number of training epochs
num_epochs = 5
# Learning rate for optimizers
lr = 0.0002
# Beta1 hyperparam for Adam optimizers
beta1 = 0.5
dataset=dset.ImageFolder(root='celeba',transform=transforms.Compose([
        transforms.Resize(image_size),
        transforms.CenterCrop(image_size),
        transforms.ToTensor(),
        transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))]))
dataloader=torch.utils.data.DataLoader(dataset,batch_size=batch_size,shuffle=True,num_workers=workers)
real_batch=next(iter(dataloader))
plt.figure(figsize=(8,8))
plt.axis('off')
plt.title('training images')
plt.imshow(np.transpose(vutils.make_grid(real_batch[0][:64],padding=2,normalize=True),(1,2,0)))
def weights_init(m):
        classname=m.__class__.__name__
        if classname.find('Conv')!=-1:
                nn.init.normal_(m.weight.data,0.0,0.02)
        elif classname.find('Batchorm')!=-1:
                nn.init.normal_(m.weight.data,1.0,0.02)
                nn.init.constant_(m.bias.data,0)
class Generator(nn.Module):
        def __init__(self):
                super(Generator,self).__init__()
                self.main=nn.Sequential(
                        nn.ConvTranspose2d( nz, ngf * 8, 4, 1, 0, bias=False),
                        nn.BatchNorm2d(ngf * 8),
                        nn.ReLU(True),
                        nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
                        nn.BatchNorm2d(ngf * 4),
                        nn.ReLU(True),
                        nn.ConvTranspose2d( ngf * 4, ngf * 2, 4, 2, 1, bias=False),
                        nn.BatchNorm2d(ngf * 2),
                        nn.ReLU(True),
                        nn.ConvTranspose2d( ngf * 2, ngf, 4, 2, 1, bias=False),
                        nn.BatchNorm2d(ngf),
                        nn.ReLU(True),
                        nn.ConvTranspose2d( ngf, nc, 4, 2, 1, bias=False),
                        nn.Tanh())
        def forward(self,input):
                return self.main(input)
class Discriminator(nn.Module):
        def __init__(self):
                super(Discriminator,self).__init__()
                self.main=nn.Sequential(
                        nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
                        nn.LeakyReLU(0.2, inplace=True),
                        nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
                        nn.BatchNorm2d(ndf * 2),
                        nn.LeakyReLU(0.2, inplace=True),
                        nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
                        nn.BatchNorm2d(ndf * 4),
                        nn.LeakyReLU(0.2, inplace=True),
                        nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
                        nn.BatchNorm2d(ndf * 8),
                        nn.LeakyReLU(0.2, inplace=True),
                        nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
                        nn.Sigmoid())
        def forward(self,input):
                return self.main(input)
netD=Discriminator()
netD.apply(weights_init)
print(netD)
netG=Generator()
netG.apply(weights_init)
print(netG)
criterion=nn.BCELoss()
fixed_noise=torch.randn(64,nz,1,1)
real_label=1
fake_label=0
optimizerD=optim.Adam(netD.parameters(),lr=lr,betas=(beta1,0.999))
optimizerG=optim.Adam(netG.parameters(),lr=lr,betas=(beta1,0.999))
img_list=[]
G_losses=[]
D_losses=[]
iters=0
for epoch in range(num_epochs):
        for i ,data in enumerate(dataloader,0):
                netD.zero_grad()
                real_cpu=data[0]
                b_size=real_cpu.size(0)
                label=torch.full((b_size,),real_label)
                output=netD(real_cpu).view(-1)
                errD_real=criterion(output,label)
                errD_real.backward()
                D_x=output.mean().item()
                noise=torch.randn(b_size,nz,1,1)
                fake=netG(noise)
                label.fill_(fake_label)
                output=netD(fake.detach()).view(-1)
                errD_fake=criterion(output,label)
                errD_fake.backward()
                D_G_z1=output.mean().item()
                errD=errD_real+errD_fake
                optimizerD.step()
                netG.zero_grad()
                label.fill_(real_label)
                output=netD(fake).view(-1)
                errG=criterion(output,label)
                errG.backward()
                D_G_z2=output.mean().item()
                optimizerG.step()
                if i % 50 == 0:
                        print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f'
                              % (epoch, num_epochs, i, len(dataloader),
                                 errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))
                G_losses.append(errG.item())
                D_losses.append(errD.item())
                if (iters % 500 == 0) or ((epoch == num_epochs-1) and (i == len(dataloader)-1)):
                        with torch.no_grad():
                                fake = netG(fixed_noise).detach().cpu()
                        img_list.append(vutils.make_grid(fake, padding=2, normalize=True))

                iters += 1
                
                