from keras.datasets import mnist
from keras.layers.merge import _Merge
from keras.layers import Input, Dense, Reshape, Flatten, Dropout
from keras.layers import BatchNormalization, Activation, ZeroPadding2D
from keras.layers.advanced_activations import LeakyReLU
from keras.layers.convolutional import UpSampling2D, Conv2D
from keras.models import Sequential, Model
from keras.optimizers import RMSprop
from functools import partial
from augmentation.discriminators import wasserstein_discriminator
from augmentation.generators import wasserstein_generator

import keras.backend as K

import matplotlib.pyplot as plt

import sys

import numpy as np

class RandomWeightedAverage(_Merge):
    """Provides a (random) weighted average between real and generated image samples"""
    def _merge_function(self, inputs):
        alpha = K.random_uniform((32, 1, 1, 1))
        return (alpha * inputs[0]) + ((1 - alpha) * inputs[1])

class WGANGP():
    def __init__(self, img_shape, noise_shape):
        self.img_shape = img_shape
        self.noise_shape = noise_shape

        # Following parameter and optimizer set as recommended in paper
        self.n_critic = 5
        optimizer = RMSprop(lr=0.00005)

        # Build the generator and critic
        self.generator = wasserstein_generator.create_model(noise_shape)
        self.critic = wasserstein_discriminator.create_model(img_shape)

        #-------------------------------
        # Construct Computational Graph
        #       for the Critic
        #-------------------------------

        # Freeze generator's layers while training critic
        self.generator.trainable = False

        # Image input (real sample)
        real_img = Input(shape=self.img_shape)

        # Noise input
        z_disc = Input(shape=(self.noise_shape,))
        # Generate image based of noise (fake sample)
        fake_img = self.generator(z_disc)

        # Discriminator determines validity of the real and fake images
        fake = self.critic(fake_img)
        valid = self.critic(real_img)

        # Construct weighted average between real and fake images
        interpolated_img = RandomWeightedAverage()([real_img, fake_img])
        # Determine validity of weighted sample
        validity_interpolated = self.critic(interpolated_img)

        # Use Python partial to provide loss function with additional
        # 'averaged_samples' argument
        partial_gp_loss = partial(self.gradient_penalty_loss,
                          averaged_samples=interpolated_img)
        partial_gp_loss.__name__ = 'gradient_penalty' # Keras requires function names

        self.critic_model = Model(inputs=[real_img, z_disc],
                            outputs=[valid, fake, validity_interpolated])
        self.critic_model.compile(loss=[self.wasserstein_loss,
                                              self.wasserstein_loss,
                                              partial_gp_loss],
                                        optimizer=optimizer,
                                        loss_weights=[1, 1, 10])
        #-------------------------------
        # Construct Computational Graph
        #         for Generator
        #-------------------------------

        # For the generator we freeze the critic's layers
        self.critic.trainable = False
        self.generator.trainable = True

        # Sampled noise for input to generator
        z_gen = Input(shape=(100,))
        # Generate images based of noise
        img = self.generator(z_gen)
        # Discriminator determines validity
        valid = self.critic(img)
        # Defines generator model
        self.generator_model = Model(z_gen, valid)
        self.generator_model.compile(loss=self.wasserstein_loss, optimizer=optimizer)


    def gradient_penalty_loss(self, y_true, y_pred, averaged_samples):
        """
        Computes gradient penalty based on prediction and weighted real / fake samples
        """
        gradients = K.gradients(y_pred, averaged_samples)[0]
        # compute the euclidean norm by squaring ...
        gradients_sqr = K.square(gradients)
        #   ... summing over the rows ...
        gradients_sqr_sum = K.sum(gradients_sqr,
                                  axis=np.arange(1, len(gradients_sqr.shape)))
        #   ... and sqrt
        gradient_l2_norm = K.sqrt(gradients_sqr_sum)
        # compute lambda * (1 - ||grad||)^2 still for each single sample
        gradient_penalty = K.square(1 - gradient_l2_norm)
        # return the mean as loss over all the batch samples
        return K.mean(gradient_penalty)


    def wasserstein_loss(self, y_true, y_pred):
        return K.mean(y_true * y_pred)

    def build_generator(self):

        model = Sequential()

        model.add(Dense(128 * 7 * 7, activation="relu", input_dim=self.noise_shape))
        model.add(Reshape((7, 7, 128)))
        model.add(UpSampling2D())
        model.add(Conv2D(128, kernel_size=4, padding="same"))
        model.add(BatchNormalization(momentum=0.8))
        model.add(Activation("relu"))
        model.add(UpSampling2D())
        model.add(Conv2D(64, kernel_size=4, padding="same"))
        model.add(BatchNormalization(momentum=0.8))
        model.add(Activation("relu"))
        model.add(Conv2D(self.channels, kernel_size=4, padding="same"))
        model.add(Activation("tanh"))

        model.summary()

        noise = Input(shape=(self.noise_shape,))
        img = model(noise)

        return Model(noise, img)

    def build_critic(self):

        model = Sequential()

        model.add(Conv2D(16, kernel_size=3, strides=2, input_shape=self.img_shape, padding="same"))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dropout(0.25))
        model.add(Conv2D(32, kernel_size=3, strides=2, padding="same"))
        model.add(ZeroPadding2D(padding=((0,1),(0,1))))
        model.add(BatchNormalization(momentum=0.8))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dropout(0.25))
        model.add(Conv2D(64, kernel_size=3, strides=2, padding="same"))
        model.add(BatchNormalization(momentum=0.8))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dropout(0.25))
        model.add(Conv2D(128, kernel_size=3, strides=1, padding="same"))
        model.add(BatchNormalization(momentum=0.8))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dropout(0.25))
        model.add(Flatten())
        model.add(Dense(1))

        model.summary()

        img = Input(shape=self.img_shape)
        validity = model(img)

        return Model(img, validity)

    def train(self, epochs, batch_size, sample_interval=50):

        # Adversarial ground truths
        valid = -np.ones((batch_size, 1))
        fake =  np.ones((batch_size, 1))
        dummy = np.zeros((batch_size, 1)) # Dummy gt for gradient penalty
        for epoch in range(epochs):

            for _ in range(self.n_critic):

                # ---------------------
                #  Train Discriminator
                # ---------------------

                # Select a random batch of images
                idx = np.random.randint(0, X_train.shape[0], batch_size)
                imgs = X_train[idx]
                # Sample generator input
                noise = np.random.normal(0, 1, (batch_size, self.noise_shape))
                # Train the critic
                d_loss = self.critic_model.train_on_batch([imgs, noise],
                                                                [valid, fake, dummy])

            # ---------------------
            #  Train Generator
            # ---------------------

            g_loss = self.generator_model.train_on_batch(noise, valid)

            # Plot the progress
            print ("%d [D loss: %f] [G loss: %f]" % (epoch, d_loss[0], g_loss))

            # If at save interval => save generated image samples
            if epoch % sample_interval == 0:
                self.sample_images(epoch)

    def sample_images(self, epoch):
        r, c = 5, 5
        noise = np.random.normal(0, 1, (r * c, self.noise_shape))
        gen_imgs = self.generator.predict(noise)

        # Rescale images 0 - 1
        gen_imgs = 0.5 * gen_imgs + 0.5

        fig, axs = plt.subplots(r, c)
        cnt = 0
        for i in range(r):
            for j in range(c):
                axs[i,j].imshow(gen_imgs[cnt, :,:,0], cmap='gray')
                axs[i,j].axis('off')
                cnt += 1
        fig.savefig("images/mnist_%d.png" % epoch)
        plt.close()

[((24, 16, 24, 47), 'keras.backend.random_uniform', 'K.random_uniform', ({(24, 33, 24, 46): '(32, 1, 1, 1)'}, {}), '((32, 1, 1, 1))', True, 'import keras.backend as K\n'), ((34, 20, 34, 39), 'keras.optimizers.RMSprop', 'RMSprop', (), '', False, 'from keras.optimizers import RMSprop\n'), ((37, 25, 37, 72), 'augmentation.generators.wasserstein_generator.create_model', 'wasserstein_generator.create_model', ({(37, 60, 37, 71): 'noise_shape'}, {}), '(noise_shape)', False, 'from augmentation.generators import wasserstein_generator\n'), ((38, 22, 38, 71), 'augmentation.discriminators.wasserstein_discriminator.create_model', 'wasserstein_discriminator.create_model', ({(38, 61, 38, 70): 'img_shape'}, {}), '(img_shape)', False, 'from augmentation.discriminators import wasserstein_discriminator\n'), ((49, 19, 49, 46), 'keras.layers.Input', 'Input', (), '', False, 'from keras.layers import Input, Dense, Reshape, Flatten, Dropout\n'), ((52, 17, 52, 49), 'keras.layers.Input', 'Input', (), '', False, 'from keras.layers import Input, Dense, Reshape, Flatten, Dropout\n'), ((67, 26, 68, 60), 'functools.partial', 'partial', (), '', False, 'from functools import partial\n'), ((71, 28, 72, 73), 'keras.models.Model', 'Model', (), '', False, 'from keras.models import Sequential, Model\n'), ((88, 16, 88, 35), 'keras.layers.Input', 'Input', (), '', False, 'from keras.layers import Input, Dense, Reshape, Flatten, Dropout\n'), ((94, 31, 94, 50), 'keras.models.Model', 'Model', ({(94, 37, 94, 42): 'z_gen', (94, 44, 94, 49): 'valid'}, {}), '(z_gen, valid)', False, 'from keras.models import Sequential, Model\n'), ((104, 24, 104, 43), 'keras.backend.square', 'K.square', ({(104, 33, 104, 42): 'gradients'}, {}), '(gradients)', True, 'import keras.backend as K\n'), ((109, 27, 109, 52), 'keras.backend.sqrt', 'K.sqrt', ({(109, 34, 109, 51): 'gradients_sqr_sum'}, {}), '(gradients_sqr_sum)', True, 'import keras.backend as K\n'), ((111, 27, 111, 57), 'keras.backend.square', 'K.square', ({(111, 36, 111, 56): '1 - gradient_l2_norm'}, {}), '(1 - gradient_l2_norm)', True, 'import keras.backend as K\n'), ((113, 15, 113, 39), 'keras.backend.mean', 'K.mean', ({(113, 22, 113, 38): 'gradient_penalty'}, {}), '(gradient_penalty)', True, 'import keras.backend as K\n'), ((117, 15, 117, 38), 'keras.backend.mean', 'K.mean', ({(117, 22, 117, 37): '(y_true * y_pred)'}, {}), '(y_true * y_pred)', True, 'import keras.backend as K\n'), ((121, 16, 121, 28), 'keras.models.Sequential', 'Sequential', ({}, {}), '()', False, 'from keras.models import Sequential, Model\n'), ((138, 16, 138, 48), 'keras.layers.Input', 'Input', (), '', False, 'from keras.layers import Input, Dense, Reshape, Flatten, Dropout\n'), ((141, 15, 141, 32), 'keras.models.Model', 'Model', ({(141, 21, 141, 26): 'noise', (141, 28, 141, 31): 'img'}, {}), '(noise, img)', False, 'from keras.models import Sequential, Model\n'), ((145, 16, 145, 28), 'keras.models.Sequential', 'Sequential', ({}, {}), '()', False, 'from keras.models import Sequential, Model\n'), ((168, 14, 168, 41), 'keras.layers.Input', 'Input', (), '', False, 'from keras.layers import Input, Dense, Reshape, Flatten, Dropout\n'), ((171, 15, 171, 35), 'keras.models.Model', 'Model', ({(171, 21, 171, 24): 'img', (171, 26, 171, 34): 'validity'}, {}), '(img, validity)', False, 'from keras.models import Sequential, Model\n'), ((177, 16, 177, 40), 'numpy.ones', 'np.ones', ({(177, 24, 177, 39): '(batch_size, 1)'}, {}), '((batch_size, 1))', True, 'import numpy as np\n'), ((178, 16, 178, 41), 'numpy.zeros', 'np.zeros', ({(178, 25, 178, 40): '(batch_size, 1)'}, {}), '((batch_size, 1))', True, 'import numpy as np\n'), ((211, 16, 211, 65), 'numpy.random.normal', 'np.random.normal', ({(211, 33, 211, 34): '0', (211, 36, 211, 37): '1', (211, 39, 211, 64): '(r * c, self.noise_shape)'}, {}), '(0, 1, (r * c, self.noise_shape))', True, 'import numpy as np\n'), ((217, 19, 217, 37), 'matplotlib.pyplot.subplots', 'plt.subplots', ({(217, 32, 217, 33): 'r', (217, 35, 217, 36): 'c'}, {}), '(r, c)', True, 'import matplotlib.pyplot as plt\n'), ((225, 8, 225, 19), 'matplotlib.pyplot.close', 'plt.close', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((102, 20, 102, 57), 'keras.backend.gradients', 'K.gradients', ({(102, 32, 102, 38): 'y_pred', (102, 40, 102, 56): 'averaged_samples'}, {}), '(y_pred, averaged_samples)', True, 'import keras.backend as K\n'), ((123, 18, 123, 83), 'keras.layers.Dense', 'Dense', (), '', False, 'from keras.layers import Input, Dense, Reshape, Flatten, Dropout\n'), ((124, 18, 124, 38), 'keras.layers.Reshape', 'Reshape', ({(124, 26, 124, 37): '(7, 7, 128)'}, {}), '((7, 7, 128))', False, 'from keras.layers import Input, Dense, Reshape, Flatten, Dropout\n'), ((125, 18, 125, 32), 'keras.layers.convolutional.UpSampling2D', 'UpSampling2D', ({}, {}), '()', False, 'from keras.layers.convolutional import UpSampling2D, Conv2D\n'), ((126, 18, 126, 60), 'keras.layers.convolutional.Conv2D', 'Conv2D', (), '', False, 'from keras.layers.convolutional import UpSampling2D, Conv2D\n'), ((127, 18, 127, 50), 'keras.layers.BatchNormalization', 'BatchNormalization', (), '', False, 'from keras.layers import BatchNormalization, Activation, ZeroPadding2D\n'), ((128, 18, 128, 36), 'keras.layers.Activation', 'Activation', ({(128, 29, 128, 35): '"""relu"""'}, {}), "('relu')", False, 'from keras.layers import BatchNormalization, Activation, ZeroPadding2D\n'), ((129, 18, 129, 32), 'keras.layers.convolutional.UpSampling2D', 'UpSampling2D', ({}, {}), '()', False, 'from keras.layers.convolutional import UpSampling2D, Conv2D\n'), ((130, 18, 130, 59), 'keras.layers.convolutional.Conv2D', 'Conv2D', (), '', False, 'from keras.layers.convolutional import UpSampling2D, Conv2D\n'), ((131, 18, 131, 50), 'keras.layers.BatchNormalization', 'BatchNormalization', (), '', False, 'from keras.layers import BatchNormalization, Activation, ZeroPadding2D\n'), ((132, 18, 132, 36), 'keras.layers.Activation', 'Activation', ({(132, 29, 132, 35): '"""relu"""'}, {}), "('relu')", False, 'from keras.layers import BatchNormalization, Activation, ZeroPadding2D\n'), ((133, 18, 133, 70), 'keras.layers.convolutional.Conv2D', 'Conv2D', (), '', False, 'from keras.layers.convolutional import UpSampling2D, Conv2D\n'), ((134, 18, 134, 36), 'keras.layers.Activation', 'Activation', ({(134, 29, 134, 35): '"""tanh"""'}, {}), "('tanh')", False, 'from keras.layers import BatchNormalization, Activation, ZeroPadding2D\n'), ((147, 18, 147, 98), 'keras.layers.convolutional.Conv2D', 'Conv2D', (), '', False, 'from keras.layers.convolutional import UpSampling2D, Conv2D\n'), ((148, 18, 148, 38), 'keras.layers.advanced_activations.LeakyReLU', 'LeakyReLU', (), '', False, 'from keras.layers.advanced_activations import LeakyReLU\n'), ((149, 18, 149, 31), 'keras.layers.Dropout', 'Dropout', ({(149, 26, 149, 30): '(0.25)'}, {}), '(0.25)', False, 'from keras.layers import Input, Dense, Reshape, Flatten, Dropout\n'), ((150, 18, 150, 70), 'keras.layers.convolutional.Conv2D', 'Conv2D', (), '', False, 'from keras.layers.convolutional import UpSampling2D, Conv2D\n'), ((151, 18, 151, 54), 'keras.layers.ZeroPadding2D', 'ZeroPadding2D', (), '', False, 'from keras.layers import BatchNormalization, Activation, ZeroPadding2D\n'), ((152, 18, 152, 50), 'keras.layers.BatchNormalization', 'BatchNormalization', (), '', False, 'from keras.layers import BatchNormalization, Activation, ZeroPadding2D\n'), ((153, 18, 153, 38), 'keras.layers.advanced_activations.LeakyReLU', 'LeakyReLU', (), '', False, 'from keras.layers.advanced_activations import LeakyReLU\n'), ((154, 18, 154, 31), 'keras.layers.Dropout', 'Dropout', ({(154, 26, 154, 30): '(0.25)'}, {}), '(0.25)', False, 'from keras.layers import Input, Dense, Reshape, Flatten, Dropout\n'), ((155, 18, 155, 70), 'keras.layers.convolutional.Conv2D', 'Conv2D', (), '', False, 'from keras.layers.convolutional import UpSampling2D, Conv2D\n'), ((156, 18, 156, 50), 'keras.layers.BatchNormalization', 'BatchNormalization', (), '', False, 'from keras.layers import BatchNormalization, Activation, ZeroPadding2D\n'), ((157, 18, 157, 38), 'keras.layers.advanced_activations.LeakyReLU', 'LeakyReLU', (), '', False, 'from keras.layers.advanced_activations import LeakyReLU\n'), ((158, 18, 158, 31), 'keras.layers.Dropout', 'Dropout', ({(158, 26, 158, 30): '(0.25)'}, {}), '(0.25)', False, 'from keras.layers import Input, Dense, Reshape, Flatten, Dropout\n'), ((159, 18, 159, 71), 'keras.layers.convolutional.Conv2D', 'Conv2D', (), '', False, 'from keras.layers.convolutional import UpSampling2D, Conv2D\n'), ((160, 18, 160, 50), 'keras.layers.BatchNormalization', 'BatchNormalization', (), '', False, 'from keras.layers import BatchNormalization, Activation, ZeroPadding2D\n'), ((161, 18, 161, 38), 'keras.layers.advanced_activations.LeakyReLU', 'LeakyReLU', (), '', False, 'from keras.layers.advanced_activations import LeakyReLU\n'), ((162, 18, 162, 31), 'keras.layers.Dropout', 'Dropout', ({(162, 26, 162, 30): '(0.25)'}, {}), '(0.25)', False, 'from keras.layers import Input, Dense, Reshape, Flatten, Dropout\n'), ((163, 18, 163, 27), 'keras.layers.Flatten', 'Flatten', ({}, {}), '()', False, 'from keras.layers import Input, Dense, Reshape, Flatten, Dropout\n'), ((164, 18, 164, 26), 'keras.layers.Dense', 'Dense', ({(164, 24, 164, 25): '(1)'}, {}), '(1)', False, 'from keras.layers import Input, Dense, Reshape, Flatten, Dropout\n'), ((176, 17, 176, 41), 'numpy.ones', 'np.ones', ({(176, 25, 176, 40): '(batch_size, 1)'}, {}), '((batch_size, 1))', True, 'import numpy as np\n'), ((188, 22, 188, 72), 'numpy.random.randint', 'np.random.randint', ({(188, 40, 188, 41): '0', (188, 43, 188, 59): 'X_train.shape[0]', (188, 61, 188, 71): 'batch_size'}, {}), '(0, X_train.shape[0], batch_size)', True, 'import numpy as np\n'), ((191, 24, 191, 78), 'numpy.random.normal', 'np.random.normal', ({(191, 41, 191, 42): '0', (191, 44, 191, 45): '1', (191, 47, 191, 77): '(batch_size, self.noise_shape)'}, {}), '(0, 1, (batch_size, self.noise_shape))', True, 'import numpy as np\n')]