# This Python 3 environment comes with many helpful analytics libraries installed
# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python
# For example, here's several helpful packages to load

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

# Input data files are available in the read-only "../input/" directory
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory

import os
for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        os.path.join(dirname, filename)

# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All" 
# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session
"""
1. **Make train dataset (named "images") using anime-faces dataset from kaggle**
2. **Define Hyperparameters**
3. **Make GAN model (Generator, Discriminator, adversarial model)**
4. **Train Discriminator and Adversarial model**
5. **Visualize images made by Generator**
"""
"""
**[1] Load anime-faces data and make them into training dataset**
"""
## Libraries needed for handling img data ##
import glob
from PIL import Image
from keras.preprocessing.image import load_img,img_to_array
import matplotlib.pyplot as plt
dataset_dir = '../input/anime-faces/data/*.png'
input_dim = (64,64,3)

count = 0
images = None
for image_path in glob.glob(dataset_dir):
    if count == 5000:
        break
    try:
        loaded_image = load_img(image_path,target_size=(input_dim[0],input_dim[1]))
        loaded_image = img_to_array(loaded_image)
        loaded_image = loaded_image / 127.5 - 1
        loaded_image = loaded_image.astype(np.float32)
        loaded_image = np.expand_dims(loaded_image,axis=0)
        if images is None:
            images = loaded_image
        else:
            images = np.concatenate([images,loaded_image],axis=0)
    except Exception as e:
        print("Error:{}".format(e))
    count += 1
        
print(images.shape)
## Visualize training anime images
random_idx = np.random.randint(0,len(images),36)

plt.figure(figsize=(12,12))
for i in range(len(random_idx)):
    plt.subplot(6,6,i+1)
    plt.xticks([])
    plt.yticks([])
    plt.grid('off')
    plt.imshow(images[random_idx[i]])
    
plt.tight_layout()
plt.show()
"""
**[2] Define Hyperparameters and parameters related to GAN model**
"""
## Libraries needed for building GAN model
from keras.layers import Input,Dense,Reshape,Flatten,Dropout,Activation,UpSampling2D,Conv2D,MaxPooling2D,LeakyReLU,BatchNormalization
from keras.models import Model
from keras.optimizers import Adam
from keras.initializers import RandomNormal
from keras import backend as K

## HyYPERPARAMETERS related to training process
# input_dim = (64,64,3)
z_dim = 512
BATCH_SIZE = 256
epochs = 220
G_lr = 0.0004
D_lr = 0.0004
weight_initializer = RandomNormal(mean=0.,stddev=0.1)

## G
initial_image_size = (4,4,128)
G_filters = [128,64,32,3]
G_kernel_sizes = [5,5,5,5]
G_strides = [1,1,1,1]

## D
D_filters = [16,32,64,128]
D_kernel_sizes = [5,5,3,3]
D_strides = [1,1,1,1]
"""
**[3] Make Generator ("G") , Discriminator ("D") , and Adversarial model (Generator connected to Discriminator) ("combined")**
"""
def build_generator():
    X = Input(shape=[z_dim])
    
    H = Dense(1024)(X)
    H = LeakyReLU(alpha=0.2)(H)
    
    H = Dense(np.prod(initial_image_size))(H)
    H = BatchNormalization()(H)
    H = LeakyReLU(alpha=0.2)(H)
    
    H = Reshape(initial_image_size)(H)
    
    for i in range(len(G_filters)):
        H = UpSampling2D(size=(2,2))(H)
        H = Conv2D(filters=G_filters[i],kernel_size=G_kernel_sizes[i],strides=G_strides[i],padding='same',kernel_initializer=weight_initializer)(H)
        
        if i < len(G_filters)-1:
            H = BatchNormalization()(H)
            H = LeakyReLU(alpha=0.2)(H)
        else:
            Y = Activation('tanh')(H)
        
    generator = Model(X,Y)
    
    return generator
def build_discriminator():
    X = Input(shape=input_dim)
    
    H = X
    for i in range(len(D_filters)):
        H = Conv2D(filters=D_filters[i],kernel_size=D_kernel_sizes[i],strides=D_strides[i],padding='same',kernel_initializer=weight_initializer)(H)
        H = LeakyReLU(alpha=0.2)(H)
        H = MaxPooling2D(pool_size=(2,2))(H)
        
    H = Flatten()(H)
    H = Dense(1024)(H)
    H = LeakyReLU(alpha=0.2)(H)
    
    H = Dense(1)(H)
    Y = Activation('sigmoid')(H)
    
    discriminator = Model(X,Y)
    
    return discriminator
def make_trainable(network,value):
    network.trainable = value
    for layer in network.layers:
        layer.trainable = value
"""
**Compile Discriminator model**
"""
G = build_generator()
D = build_discriminator()

D_opt = Adam(lr=D_lr)
D.compile(loss='binary_crossentropy',optimizer=D_opt)
D.summary()
"""
**Compile Adversarial model**
"""
combined_opt = Adam(lr=G_lr)

## first, freeze Discriminator network
make_trainable(D,False)

Z = Input(shape=[z_dim])
generated = G(Z)
valid = D(generated)

combined = Model(Z,valid)
combined.compile(loss='binary_crossentropy',optimizer=combined_opt)

## unfreeze Discriminator because we have to train Discriminator too
make_trainable(D,True)

combined.summary()
## Functions related to Visualizing

def denormalize(img):
    img = (img+1)*127.5
    return img.astype(np.uint8)

def show_generated():
    noise = np.random.normal(0,1,size=(16,z_dim))
    generated = G.predict(noise)
    
    plt.figure(figsize=(12,12))
    for i in range(16):
        plt.subplot(4,4,i+1)
        plt.grid(False)
        plt.xticks([])
        plt.yticks([])
        ##plt.imshow(denormalize(generated[i]))
        plt.imshow(generated[i])
    plt.tight_layout()
    plt.show()
"""
**[4] Train Discriminator and Generator**
"""
# dictionary for recording losses
losses = {"d":[],"g":[]}

for epoch in range(epochs):

    steps = int(images.shape[0]/BATCH_SIZE)
    for step in range(steps):
        z_noise = np.random.normal(0,1,size=(BATCH_SIZE,z_dim))
        fake_img = G.predict_on_batch(z_noise)
        
        ## Train D ##
        real_img = images[step*BATCH_SIZE:(step+1)*BATCH_SIZE]
        ones = np.ones((BATCH_SIZE,))
        zeros = np.zeros((BATCH_SIZE,))
        
        D_loss_real = D.train_on_batch(real_img,ones)
        D_loss_fake = D.train_on_batch(fake_img,zeros)
        D_loss = (D_loss_real + D_loss_fake) / 2
        
        losses["d"].append(D_loss)

        ## Train Combined ##
        G_loss = combined.train_on_batch(z_noise,ones)
        losses["g"].append(G_loss)
        
    print("Epoch %d     [D_LOSS : %.8f]    [G_LOSS : %.8f]"%(epoch+1,D_loss,G_loss))
    if (epoch+1) % 25 ==0:
        show_generated()
"""
**Plot of Loss Fluctuations**
"""
plt.figure(figsize=(8,12))

plt.subplot(3,1,1)
plt.title("Discriminator Loss")
plt.plot(range(len(losses["d"])),losses["d"],c="b")
plt.ylim(0,1)

plt.subplot(3,1,2)
plt.title("Generator Loss")
plt.plot(range(len(losses["g"])),losses["g"],c="r")
plt.ylim(0,1)

plt.subplot(3,1,3)
plt.title("Combined Plot")
plt.plot(range(len(losses["d"])),losses["d"],c="b",label="D_Loss")
plt.plot(range(len(losses["g"])),losses["g"],c="r",label="G_Loss")
plt.legend()
plt.ylim(0,1)

plt.tight_layout()
plt.show()
show_generated()
noise = np.random.normal(0,1,size=(16,z_dim))
generated = G.predict(noise)
    
plt.figure(figsize=(10,10))
for i in range(16):
    plt.subplot(4,4,i+1)
    plt.grid(False)
    plt.xticks([])
    plt.yticks([])
    #plt.imshow(generated[i])
    plt.imshow(denormalize(generated[i]))
plt.tight_layout()
plt.show()
noise = np.random.normal(0,1,size=(16,z_dim))
generated = G.predict(noise)
    
plt.figure(figsize=(10,10))
for i in range(16):
    plt.subplot(4,4,i+1)
    plt.grid(False)
    plt.xticks([])
    plt.yticks([])
    plt.imshow(generated[i],cmap='binary')
plt.tight_layout()
plt.show()