%matplotlib inline
import importlib
import utils2; importlib.reload(utils2)
from utils2 import *

from tqdm import tqdm

"""
Explanation: Generative Adversarial Networks in Keras
End of explanation
"""


from keras.datasets import mnist
(X_train, y_train), (X_test, y_test) = mnist.load_data()
X_train.shape

n = len(X_train)

X_train = X_train.reshape(n, -1).astype(np.float32)
X_test = X_test.reshape(len(X_test), -1).astype(np.float32)

X_train /= 255.; X_test /= 255.

"""
Explanation: The original GAN!
See this paper for details of the approach we'll try first for our first GAN. We'll see if we can generate hand-drawn numbers based on MNIST, so let's load that dataset first.
We'll be refering to the discriminator as 'D' and the generator as 'G'.
End of explanation
"""


def plot_gen(G, n_ex=16):
    plot_multi(G.predict(noise(n_ex)).reshape(n_ex, 28,28), cmap='gray')

"""
Explanation: Train
This is just a helper to plot a bunch of generated images.
End of explanation
"""


def noise(bs): return np.random.rand(bs,100)

"""
Explanation: Create some random data for the generator.
End of explanation
"""


def data_D(sz, G):
    real_img = X_train[np.random.randint(0,n,size=sz)]
    X = np.concatenate((real_img, G.predict(noise(sz))))
    return X, [0]*sz + [1]*sz  # D should predict 1 for data generated by G, 0 for real data

def make_trainable(net, val):
    net.trainable = val
    for l in net.layers: l.trainable = val

"""
Explanation: Create a batch of some real and some generated data, with appropriate labels, for the discriminator.
End of explanation
"""


def train(D, G, m, nb_epoch=5000, bs=128):
    dl,gl=[],[]
    for e in tqdm(range(nb_epoch)):
        X,y = data_D(bs//2, G)
        dl.append(D.train_on_batch(X,y))
        make_trainable(D, False)
        gl.append(m.train_on_batch(noise(bs), np.zeros([bs])))
        make_trainable(D, True)
    return dl,gl

"""
Explanation: Train a few epochs, and return the losses for D and G. In each epoch we:

Train D on one batch from data_D()
Train G to create images that the discriminator predicts as real.
End of explanation
"""


MLP_G = Sequential([
    Dense(200, input_shape=(100,), activation='relu'),
    Dense(400, activation='relu'),
    Dense(784, activation='sigmoid'),
])

MLP_D = Sequential([
    Dense(300, input_shape=(784,), activation='relu'),
    Dense(300, activation='relu'),
    Dense(1, activation='sigmoid'),
])
MLP_D.compile(Adam(1e-4), "binary_crossentropy")

MLP_m = Sequential([MLP_G,MLP_D])
MLP_m.compile(Adam(1e-4), "binary_crossentropy")

dl,gl = train(MLP_D, MLP_G, MLP_m, 8000)

"""
Explanation: MLP GAN
We'll keep thinks simple by making D & G plain ole' MLPs.
End of explanation
"""


plt.plot(dl[100:])

plt.plot(gl[100:])

"""
Explanation: The loss plots for most GANs are nearly impossible to interpret - which is one of the things that make them hard to train.
End of explanation
"""


plot_gen(MLP_G)

"""
Explanation: This is what's known in the literature as "mode collapse".
End of explanation
"""


X_train = X_train.reshape(n, 28, 28, 1)
X_test = X_test.reshape(len(X_test), 28, 28, 1)

"""
Explanation: OK, so that didn't work. Can we do better?...
DCGAN
There's lots of ideas out there to make GANs train better, since they are notoriously painful to get working. The paper introducing DCGANs is the main basis for our next section. Add see https://github.com/soumith/ganhacks for many tips!
Because we're using a CNN from now on, we'll reshape our digits into proper images.
End of explanation
"""


# Keras 2
CNN_G = Sequential([
    Dense(512*7*7, input_dim=100),
    advanced_activations.LeakyReLU(),  # alpha=0.3 default
    BatchNormalization(),
    Reshape((7, 7, 512)),
    UpSampling2D(),
    Conv2D(64, (3, 3), padding='same'),
    advanced_activations.LeakyReLU(),
    BatchNormalization(),
    UpSampling2D(),
    Conv2D(32, (3, 3), padding='same'),
    advanced_activations.LeakyReLU(),
    BatchNormalization(),
    Conv2D(1, (1, 1), padding='same', activation='sigmoid')
])

"""
Explanation: Our generator uses a number of upsampling steps as suggested in the above papers. We use nearest neighbor upsampling rather than fractionally strided convolutions, as discussed in our style transfer notebook.
End of explanation
"""


# Keras 2
CNN_D = Sequential([
    Conv2D(256, (5, 5), strides=(2,2), padding='same', 
                  input_shape=(28, 28, 1)),
    advanced_activations.LeakyReLU(),
    Conv2D(512, (5, 5), strides=(2,2), padding='same'),
    advanced_activations.LeakyReLU(),
    Flatten(),
    Dense(256),
    advanced_activations.LeakyReLU(),
    Dense(1, activation = 'sigmoid')
])

CNN_D.compile(Adam(1e-3), "binary_crossentropy")

"""
Explanation: The discriminator uses a few downsampling steps through strided convolutions.
End of explanation
"""


sz = n//200
x1 = np.concatenate([np.random.permutation(X_train)[:sz], CNN_G.predict(noise(sz))])
CNN_D.fit(x1, [0]*sz + [1]*sz, batch_size=128, epochs=1, verbose=2)

CNN_m = Sequential([CNN_G, CNN_D])
CNN_m.compile(Adam(1e-4), "binary_crossentropy")

K.set_value(CNN_D.optimizer.lr, 1e-3)
K.set_value(CNN_m.optimizer.lr, 1e-3)

"""
Explanation: We train D a "little bit" so it can at least tell a real image from random noise.
End of explanation
"""


dl,gl = train(CNN_D, CNN_G, CNN_m, 2500)

plt.plot(dl[10:])

plt.plot(gl[10:])

"""
Explanation: Now we can train D & G iteratively.
End of explanation
"""


plot_gen(CNN_G)

"""
Explanation: Better than our first effort, but still a lot to be desired:...
End of explanation
"""


