import os

import numpy as np
import tensorflow.keras as keras
from tensorflow.keras.preprocessing import image
from tensorflow.keras import layers

latent_dim = 32
height = 32
width = 32
channels = 3

# =========================================

generator_input = keras.Input(shape=(latent_dim,))

# First, transform the input into a 16x16 128-channels feature map
x = layers.Dense(128 * 16 * 16)(generator_input)
x = layers.LeakyReLU()(x)
x = layers.Reshape((16, 16, 128))(x)

# Then, add a convolution layer
x = layers.Conv2D(256, 5, padding='same')(x)
x = layers.LeakyReLU()(x)

# Upsample to 32x32
x = layers.Conv2DTranspose(256, 4, strides=2, padding='same')(x)
x = layers.LeakyReLU()(x)

# Few more conv layers
x = layers.Conv2D(256, 5, padding='same')(x)
x = layers.LeakyReLU()(x)
x = layers.Conv2D(256, 5, padding='same')(x)
x = layers.LeakyReLU()(x)

# Produce a 32x32 1-channel feature map
x = layers.Conv2D(channels, 7, activation='tanh', padding='same')(x)
generator = keras.models.Model(generator_input, x)
generator.summary()

# =========================================

discriminator_input = layers.Input(shape=(height, width, channels))
x = layers.Conv2D(128, 3)(discriminator_input)
x = layers.LeakyReLU()(x)
x = layers.Conv2D(128, 4, strides=2)(x)
x = layers.LeakyReLU()(x)
x = layers.Conv2D(128, 4, strides=2)(x)
x = layers.LeakyReLU()(x)
x = layers.Conv2D(128, 4, strides=2)(x)
x = layers.LeakyReLU()(x)
x = layers.Flatten()(x)

# One dropout layer - important trick!
x = layers.Dropout(0.4)(x)

# Classification layer
x = layers.Dense(1, activation='sigmoid')(x)

discriminator = keras.models.Model(discriminator_input, x)
discriminator.summary()

# To stabilize training, we use learning rate decay
# and gradient clipping (by value) in the optimizer.
discriminator_optimizer = keras.optimizers.RMSprop(lr=0.0008, clipvalue=1.0, decay=1e-8)
discriminator.compile(optimizer=discriminator_optimizer, loss='binary_crossentropy')

# =========================================

# Set discriminator weights to non-trainable
# (will only apply to the `gan` model)
discriminator.trainable = False

gan_input = keras.Input(shape=(latent_dim,))
gan_output = discriminator(generator(gan_input))
gan = keras.models.Model(gan_input, gan_output)

gan_optimizer = keras.optimizers.RMSprop(lr=0.0004, clipvalue=1.0, decay=1e-8)
gan.compile(optimizer=gan_optimizer, loss='binary_crossentropy')

gan.summary()

# =========================================

# Load CIFAR10 data
(x_train, y_train), (_, _) = keras.datasets.cifar10.load_data()

# Select frog images (class 6)
x_train = x_train[y_train.flatten() == 6]

# Normalize data
x_train = x_train.reshape(
    (x_train.shape[0],) + (height, width, channels)).astype('float32') / 255.

iterations = 10000
batch_size = 20
save_dir = './gan_images'

# Start training loop
start = 0
for step in range(iterations):
    # Sample random points in the latent space
    random_latent_vectors = np.random.normal(size=(batch_size, latent_dim))

    # Decode them to fake images
    generated_images = generator.predict(random_latent_vectors)

    # Combine them with real images
    stop = start + batch_size
    real_images = x_train[start: stop]
    combined_images = np.concatenate([generated_images, real_images])

    # Assemble labels discriminating real from fake images
    labels = np.concatenate([np.ones((batch_size, 1)),
                             np.zeros((batch_size, 1))])
    # Add random noise to the labels - important trick!
    labels += 0.05 * np.random.random(labels.shape)

    # Train the discriminator
    d_loss = discriminator.train_on_batch(combined_images, labels)

    # sample random points in the latent space
    random_latent_vectors = np.random.normal(size=(batch_size, latent_dim))

    # Assemble labels that say "all real images"
    misleading_targets = np.zeros((batch_size, 1))

    # Train the generator (via the gan model,
    # where the discriminator weights are frozen)
    a_loss = gan.train_on_batch(random_latent_vectors, misleading_targets)

    start += batch_size
    if start > len(x_train) - batch_size:
        start = 0

    # Occasionally save / plot
    if step % 100 == 0:
        # Save model weights
        gan.save_weights('gan.h5')

        # Print metrics
        print('discriminator loss at step %s: %s' % (step, d_loss))
        print('adversarial loss at step %s: %s' % (step, a_loss))

        # Save one generated image
        img = image.array_to_img(generated_images[0] * 255., scale=False)
        img.save(os.path.join(save_dir, 'generated_frog' + str(step) + '.png'))

        # Save one real image, for comparison
        img = image.array_to_img(real_images[0] * 255., scale=False)
        img.save(os.path.join(save_dir, 'real_frog' + str(step) + '.png'))

tensorflow.keras.layers.Dropout tensorflow.keras.layers.LeakyReLU tensorflow.keras.Input numpy.random.random tensorflow.keras.models.Model tensorflow.keras.layers.Dense tensorflow.keras.layers.Conv2DTranspose tensorflow.keras.optimizers.RMSprop tensorflow.keras.layers.Conv2D tensorflow.keras.datasets.cifar10.load_data numpy.ones numpy.concatenate numpy.random.normal tensorflow.keras.layers.Reshape tensorflow.keras.preprocessing.image.array_to_img tensorflow.keras.layers.Flatten numpy.zeros tensorflow.keras.layers.Input
[(15, 'tensorflow.keras.Input', 'keras.Input', ([], {'shape': '(latent_dim,)'}), True, 'import tensorflow.keras as keras\n'), (38, 'tensorflow.keras.models.Model', 'keras.models.Model', (['generator_input', 'x'], {}), True, 'import tensorflow.keras as keras\n'), (43, 'tensorflow.keras.layers.Input', 'layers.Input', ([], {'shape': '(height, width, channels)'}), False, 'from tensorflow.keras import layers\n'), (60, 'tensorflow.keras.models.Model', 'keras.models.Model', (['discriminator_input', 'x'], {}), True, 'import tensorflow.keras as keras\n'), (65, 'tensorflow.keras.optimizers.RMSprop', 'keras.optimizers.RMSprop', ([], {'lr': '(0.0008)', 'clipvalue': '(1.0)', 'decay': '(1e-08)'}), True, 'import tensorflow.keras as keras\n'), (74, 'tensorflow.keras.Input', 'keras.Input', ([], {'shape': '(latent_dim,)'}), True, 'import tensorflow.keras as keras\n'), (76, 'tensorflow.keras.models.Model', 'keras.models.Model', (['gan_input', 'gan_output'], {}), True, 'import tensorflow.keras as keras\n'), (78, 'tensorflow.keras.optimizers.RMSprop', 'keras.optimizers.RMSprop', ([], {'lr': '(0.0004)', 'clipvalue': '(1.0)', 'decay': '(1e-08)'}), True, 'import tensorflow.keras as keras\n'), (86, 'tensorflow.keras.datasets.cifar10.load_data', 'keras.datasets.cifar10.load_data', ([], {}), True, 'import tensorflow.keras as keras\n'), (18, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(128 * 16 * 16)'], {}), False, 'from tensorflow.keras import layers\n'), (19, 'tensorflow.keras.layers.LeakyReLU', 'layers.LeakyReLU', ([], {}), False, 'from tensorflow.keras import layers\n'), (20, 'tensorflow.keras.layers.Reshape', 'layers.Reshape', (['(16, 16, 128)'], {}), False, 'from tensorflow.keras import layers\n'), (23, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['(256)', '(5)'], {'padding': '"""same"""'}), False, 'from tensorflow.keras import layers\n'), (24, 'tensorflow.keras.layers.LeakyReLU', 'layers.LeakyReLU', ([], {}), False, 'from tensorflow.keras import layers\n'), (27, 'tensorflow.keras.layers.Conv2DTranspose', 'layers.Conv2DTranspose', (['(256)', '(4)'], {'strides': '(2)', 'padding': '"""same"""'}), False, 'from tensorflow.keras import layers\n'), (28, 'tensorflow.keras.layers.LeakyReLU', 'layers.LeakyReLU', ([], {}), False, 'from tensorflow.keras import layers\n'), (31, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['(256)', '(5)'], {'padding': '"""same"""'}), False, 'from tensorflow.keras import layers\n'), (32, 'tensorflow.keras.layers.LeakyReLU', 'layers.LeakyReLU', ([], {}), False, 'from tensorflow.keras import layers\n'), (33, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['(256)', '(5)'], {'padding': '"""same"""'}), False, 'from tensorflow.keras import layers\n'), (34, 'tensorflow.keras.layers.LeakyReLU', 'layers.LeakyReLU', ([], {}), False, 'from tensorflow.keras import layers\n'), (37, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['channels', '(7)'], {'activation': '"""tanh"""', 'padding': '"""same"""'}), False, 'from tensorflow.keras import layers\n'), (44, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['(128)', '(3)'], {}), False, 'from tensorflow.keras import layers\n'), (45, 'tensorflow.keras.layers.LeakyReLU', 'layers.LeakyReLU', ([], {}), False, 'from tensorflow.keras import layers\n'), (46, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['(128)', '(4)'], {'strides': '(2)'}), False, 'from tensorflow.keras import layers\n'), (47, 'tensorflow.keras.layers.LeakyReLU', 'layers.LeakyReLU', ([], {}), False, 'from tensorflow.keras import layers\n'), (48, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['(128)', '(4)'], {'strides': '(2)'}), False, 'from tensorflow.keras import layers\n'), (49, 'tensorflow.keras.layers.LeakyReLU', 'layers.LeakyReLU', ([], {}), False, 'from tensorflow.keras import layers\n'), (50, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['(128)', '(4)'], {'strides': '(2)'}), False, 'from tensorflow.keras import layers\n'), (51, 'tensorflow.keras.layers.LeakyReLU', 'layers.LeakyReLU', ([], {}), False, 'from tensorflow.keras import layers\n'), (52, 'tensorflow.keras.layers.Flatten', 'layers.Flatten', ([], {}), False, 'from tensorflow.keras import layers\n'), (55, 'tensorflow.keras.layers.Dropout', 'layers.Dropout', (['(0.4)'], {}), False, 'from tensorflow.keras import layers\n'), (58, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(1)'], {'activation': '"""sigmoid"""'}), False, 'from tensorflow.keras import layers\n'), (103, 'numpy.random.normal', 'np.random.normal', ([], {'size': '(batch_size, latent_dim)'}), True, 'import numpy as np\n'), (111, 'numpy.concatenate', 'np.concatenate', (['[generated_images, real_images]'], {}), True, 'import numpy as np\n'), (123, 'numpy.random.normal', 'np.random.normal', ([], {'size': '(batch_size, latent_dim)'}), True, 'import numpy as np\n'), (126, 'numpy.zeros', 'np.zeros', (['(batch_size, 1)'], {}), True, 'import numpy as np\n'), (117, 'numpy.random.random', 'np.random.random', (['labels.shape'], {}), True, 'import numpy as np\n'), (146, 'tensorflow.keras.preprocessing.image.array_to_img', 'image.array_to_img', (['(generated_images[0] * 255.0)'], {'scale': '(False)'}), False, 'from tensorflow.keras.preprocessing import image\n'), (150, 'tensorflow.keras.preprocessing.image.array_to_img', 'image.array_to_img', (['(real_images[0] * 255.0)'], {'scale': '(False)'}), False, 'from tensorflow.keras.preprocessing import image\n'), (114, 'numpy.ones', 'np.ones', (['(batch_size, 1)'], {}), True, 'import numpy as np\n'), (115, 'numpy.zeros', 'np.zeros', (['(batch_size, 1)'], {}), True, 'import numpy as np\n')]