#!/usr/bin/env python
# coding: utf-8

# In[1]:


from __future__ import print_function
#%matplotlib inline
import os
import random
import torch
import torch.nn as nn
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
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import argparse
from IPython.display import HTML
#from Model import *
import Model as GAN_model
# Set random seed for reproducibility
manualSeed = 999
#manualSeed = random.randint(1, 10000) # use if you want new results
print("Random Seed: ", manualSeed)
random.seed(manualSeed)
torch.manual_seed(manualSeed)

ngpu = 1
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
print(device)


# In[2]:


#這是給CMD用的，我們自己令一個arg_parser class
# def common_arg_parser():
#     parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
#     parser.add_argument('--dataroot', default='/home/jeff/mount/Homework3/celeba/', type=str)
#     parser.add_argument('--batch_size', default=128, type=int)
#     parser.add_argument('--image_size', default=64, type=int)
#     parser.add_argument('--num_epochs', default=5, type=int)
#     parser.add_argument('--lr', default=0.0002, type=float)
    

#     return parser


# In[3]:


#存超參數用的，要改就直接改裡面，不過這邊沿用助教的
class arg_parser:
    def __init__(self):
        self.dataroot ='/home/jeff/mount/Homework3/celeba/'
        self.batch_size = 128
        self.image_size = 64
        self.num_epochs = 5
        self.lr = 0.0002
        #Adam要用的超參數
        self.beta1 = 0.5
args = arg_parser()


# In[4]:


def train(dataloader, generator, discriminator, optimizer_g, optimizer_d, criterion, num_epochs):
    #這兩個label要令在裡面，如果有必要的話再把它加成train function的參數之一
    real_label = 1
    fake_label = 0
    #存outcome
    img_list = []
    G_losses = []
    D_losses = []
    iters = 0
    print("Starting Training Loop...")
    # Each epoch, we have to go through every data in dataset
    for epoch in range(num_epochs):
        # Each iteration, we will get a batch data for training
        for i, data in enumerate(dataloader, 0):
            # initialize gradient for network
            discriminator.zero_grad()
            # send the data into device for computation
            real_gpu = data[0].to(device)
            b_size = real_gpu.size(0)
            
            # Send data to discriminator and calculate the loss and gradient
            # For calculate loss, you need to create label for your data
            label = torch.full((b_size,), real_label, device = device)
            output = discriminator(real_gpu).view(-1)
            
            errD_real = criterion(output, label)
            errD_real.backward()
            D_x = output.mean().item()
            
            ## Using Fake data, other steps are the same.
            noise = torch.randn(b_size, 100, 1, 1, device=device)
            fake = generator(noise)
            label.fill_(fake_label)
            output = discriminator(fake.detach()).view(-1)

            errD_fake = criterion(output, label)
            errD_fake.backward()

            D_G_z1 = output.mean().item()
            errD = errD_real + errD_fake
            optimizer_d.step()

            # Update your network
            generator.zero_grad()
            label.fill_(real_label)
            output = discriminator(fake).view(-1)
            errG = criterion(output, label)
            errG.backward()
            D_G_z2 = output.mean().item()
            optimizer_g.step()
            
 
            # Use this function to output training procedure while training
            # You can also use this function to save models and samples after fixed number of iteration
            if i % 50 == 0:
                print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f' % (epoch+1, num_epochs, i, len(dataloader), errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))
            
            # Record your loss every iteration for visualization
            G_losses.append(errG.item())
            D_losses.append(errD.item())
            
            # Check how the generator is doing by saving G's output on fixed_noise
            if (iters % 500 == 0) or ((epoch == num_epochs - 1) and (i == len(dataloader) - 1)):
                with torch.no_grad():
                    fixed_noise = torch.randn(64, 100, 1, 1, device=device)
                    fake = generator(fixed_noise).detach().cpu()
                    
                # restore images generated by GAN
                img_list.append(vutils.make_grid(fake, padding=2, normalize=True))

            iters += 1
            
            # 在所有batch結束之後把train過一遍的generator跟discrimeinator存起來
            torch.save(generator, 'generator.pkl')
            torch.save(discriminator, 'discriminator.pkl')
            
    return G_losses, D_losses, img_list


# In[5]:


def weights_init(m):
    classname = m.__class__.__name__
    print('classname:', classname)

    if classname.find('Conv') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)


# In[7]:


# Preprocessing images
dataset = dset.ImageFolder(root=args.dataroot, 
                           transform=transforms.Compose([transforms.Resize(args.image_size),
                                                         transforms.CenterCrop(args.image_size),
                                                         transforms.ToTensor(), 
                                                         transforms.Normalize((0, 0, 0), (1, 1, 1))]))
# Create the dataloader
dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True)

# Create the Discriminator
generator = GAN_model.Generator(ngpu).to(device)
# Handle multi-gpu if desired
if (device.type == 'cuda') and (ngpu > 1):
    generator = nn.DataParallel(generator, list(range(ngpu)))
generator.apply(weights_init)
# Create the Generator
discriminator = GAN_model.Discriminator(ngpu).to(device)
# Handle multi-gpu if desired
if (device.type == 'cuda') and (ngpu > 1):
    discriminator = nn.DataParallel(discriminator, list(range(ngpu)))
discriminator.apply(weights_init)
print(generator)
print(discriminator)

# Initialize BCELoss function
criterion = nn.BCELoss()

# Create batch of latent vectors that we will use to visualize
#  the progression of the generator
fixed_noise = torch.randn(64, 100, 1, 1, device=device)
# Setup Adam optimizers for both G and D
optimizerD = optim.Adam(discriminator.parameters(), lr=args.lr, betas=(args.beta1, 0.999))
optimizerG = optim.Adam(generator.parameters(), lr=args.lr, betas=(args.beta1, 0.999))

criterion = nn.BCELoss()


# In[9]:


G_losses, D_losses, images_GAN = train(dataloader, generator, discriminator, optimizerG, optimizerD, criterion, args.num_epochs)


# In[30]:


import pandas as pd
plt.figure(figsize=(14, 5))
plt.title("Generator and Discriminator Loss")
plt.plot(G_losses, label="Generator")
plt.plot(pd.Series(G_losses).rolling(100).mean(), 'b')
plt.plot(D_losses, label="Discriminator")
plt.plot(pd.Series(D_losses).rolling(100).mean(), 'b')
plt.xlabel("iterations")
plt.ylabel("Loss")
plt.legend()
plt.show()


# In[22]:


# Grab a batch of real images from the dataloader
real_batch = next(iter(dataloader))
# Plot the real images
plt.figure(figsize=(15, 15))
plt.subplot(1, 2, 1)
plt.axis("off")
plt.title("Real Images")
plt.imshow(np.transpose(vutils.make_grid(real_batch[0].to(device)[:64], padding=5, normalize=True).cpu(), (1, 2, 0)))

# Plot the fake images from the last epoch
plt.subplot(1, 2, 2)
plt.axis("off")
plt.title("Fake Images")
plt.imshow(np.transpose(images_GAN[-1], (1, 2, 0)))
plt.show()


# In[ ]:




