# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---
# # Deep Convolutional GAN (DC-GAN)
#
# Originally proposed by [Radford et al.](https://arxiv.org/pdf/1511.06434.pdf) is their work titled Unsupervised Representation Learning With Deep Convolutions Generative Adversarial Networks. This network uses a basic implementation where generator and discriminator models use convolutional layers, batch normalization and Upsampling.
# This notebook trains both networks using ADAM optimizer to play the minimax game. We showcase the effectiveness using MNIST digit generation
# [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/PacktPublishing/Hands-On-Generative-AI-with-Python-and-TensorFlow-2/blob/master/Chapter_6/deep_convolutional_gan.ipynb)
# ## Load Libraries
from tensorflow.keras import Model
from tensorflow.keras.layers import Input
from tensorflow.keras.optimizers import Adam
from tensorflow.keras import datasets
import numpy as np
# ## Load Utility Functions
from gan_utils import build_dc_discriminator
from gan_utils import build_dc_generator
from gan_utils import sample_images

# ## DCGAN Training Loop
# - As proposed in the original paper
# - Train discriminator using a mix of fake and real samples
# - Calculate discriminator loss
# - Fix the discriminator and train generator
#
# _Note : This is same as Vanilla GAN training_
def train(generator=None,discriminator=None,gan_model=None,
          epochs=1000, batch_size=128, sample_interval=50,
          z_dim=100):
    # Load MNIST train samples
    (X_train, _), (_, _) = datasets.mnist.load_data()
    # Rescale -1 to 1
    X_train = X_train / 127.5 - 1
    X_train = np.expand_dims(X_train, axis=3)
    # Prepare GAN output labels
    real_y = np.ones((batch_size, 1))
    fake_y = np.zeros((batch_size, 1))
    for epoch in range(epochs):
        # train disriminator
        # pick random real samples from X_train
        idx = np.random.randint(0, X_train.shape[0], batch_size)
        real_imgs = X_train[idx]
        # pick random noise samples (z) from a normal distribution
        noise = np.random.normal(0, 1, (batch_size, z_dim))
        # use generator model to generate output samples
        fake_imgs = generator.predict(noise)
        # calculate discriminator loss on real samples
        disc_loss_real = discriminator.train_on_batch(real_imgs, real_y)
        
        # calculate discriminator loss on fake samples
        disc_loss_fake = discriminator.train_on_batch(fake_imgs, fake_y)
        
        # overall discriminator loss
        discriminator_loss = 0.5 * np.add(disc_loss_real, disc_loss_fake)
        
        #train generator
        # pick random noise samples (z) from a normal distribution
        noise = np.random.normal(0, 1, (batch_size, z_dim))
        # use trained discriminator to improve generator
        gen_loss = gan_model.train_on_batch(noise, real_y)
        # training updates
        print ("%d [Discriminator loss: %f, acc.: %.2f%%] [Generator loss: %f]" % (epoch, 
                                                                                   discriminator_loss[0], 
                                                                                   100*discriminator_loss[1], 
                                                                                   gen_loss))
        # If at save interval => save generated image samples
        if epoch % sample_interval == 0:
            sample_images(epoch,generator)

# ## Prepare Discriminator Model
discriminator = build_dc_discriminator()
discriminator.compile(loss='binary_crossentropy',
                      optimizer=Adam(0.0001, 0.2),
                      metrics=['accuracy'])
# ## Prepare Generator Model
generator=build_dc_generator()
# ## Prepare GAN Model
# +
# Noise for generator
z_dim = 100
z = Input(shape=(z_dim,))
img = generator(z)
# Fix the discriminator
discriminator.trainable = False
# Get discriminator output
validity = discriminator(img)
# Stack discriminator on top of generator
gan_model = Model(z, validity)
gan_model.compile(loss='binary_crossentropy', optimizer=Adam(0.0005, 0.5))
gan_model.summary()
# -
# ## Train DCGAN
train(generator,discriminator,gan_model,epochs=2500, batch_size=64, sample_interval=100)
# ## Output
# Samples generated after 2500 epochs
# <img src="outputs/dc_gan_output.png">