text: '''Trains LSGAN on MNIST using Keras
  LSGAN is similar to DCGAN except for the MSE loss used by the
  Discriminator and Adversarial networks.
  
  [1] Radford, Alec, Luke Metz, and Soumith Chintala.
  "Unsupervised representation learning with deep convolutional
  generative adversarial networks." arXiv preprint arXiv:1511.06434 (2015).
  [2] Mao, Xudong, et al. "Least squares generative adversarial networks."
  2017 IEEE International Conference on Computer Vision (ICCV). IEEE, 2017.
  '''
  
  
  from tensorflow.keras.layers import Activation, Dense, Input
  from tensorflow.keras.layers import Conv2D, Flatten
  from tensorflow.keras.layers import Reshape, Conv2DTranspose
  from tensorflow.keras.layers import LeakyReLU
  from tensorflow.keras.layers import BatchNormalization
  from tensorflow.keras.layers import concatenate
  from tensorflow.keras.models import Model
  from tensorflow.keras.optimizers import RMSprop
  from tensorflow.keras.datasets import mnist
  from tensorflow.keras import backend as K
  from tensorflow.keras.models import load_model
  
  import math
  import matplotlib.pyplot as plt
  import numpy as np
  import argparse
  
  def build_generator(inputs,
                image_size,
                activation='sigmoid',
                labels=None,
                codes=None):
      """Build a Generator Model
      Stack of BN-ReLU-Conv2DTranpose to generate fake images.
      Output activation is sigmoid instead of tanh in [1].
      Sigmoid converges easily.
      Arguments:
          inputs (Layer): Input layer of the generator (the z-vector)
          image_size (int): Target size of one side
              (assuming square image)
          activation (string): Name of output activation layer
          labels (tensor): Input labels
          codes (list): 2-dim disentangled codes for InfoGAN
      Returns:
          Model: Generator Model
      """
      image_resize = image_size // 4
      # network parameters
      kernel_size = 5
      layer_filters = [128, 64, 32, 1]
  
      if labels is not None:
          if codes is None:
              # ACGAN labels
              # concatenate z noise vector and one-hot labels
              inputs = [inputs, labels]
          else:
              # infoGAN codes
              # concatenate z noise vector,
              # one-hot labels and codes 1 & 2
              inputs = [inputs, labels] + codes
          x = concatenate(inputs, axis=1)
      elif codes is not None:
          # generator 0 of StackedGAN
          inputs = [inputs, codes]
          x = concatenate(inputs, axis=1)
      else:
          # default input is just 100-dim noise (z-code)
          x = inputs
  
      x = Dense(image_resize * image_resize * layer_filters[0])(x)
      x = Reshape((image_resize, image_resize, layer_filters[0]))(x)
  
      for filters in layer_filters:
          # first two convolution layers use strides = 2
          # the last two use strides = 1
          if filters > layer_filters[-2]:
              strides = 2
          else:
              strides = 1
          x = BatchNormalization()(x)
          x = Activation('relu')(x)
          x = Conv2DTranspose(filters=filters,
                              kernel_size=kernel_size,
                              strides=strides,
                              padding='same')(x)
  
      if activation is not None:
          x = Activation(activation)(x)
  
      # generator output is the synthesized image x
      return Model(inputs, x, name='generator')
  
  
  def build_discriminator(inputs,
                    activation='sigmoid',
                    num_labels=None,
                    num_codes=None):
      """Build a Discriminator Model
      Stack of LeakyReLU-Conv2D to discriminate real from fake
      The network does not converge with BN so it is not used here
      unlike in [1]
      Arguments:
          inputs (Layer): Input layer of the discriminator (the image)
          activation (string): Name of output activation layer
          num_labels (int): Dimension of one-hot labels for ACGAN & InfoGAN
          num_codes (int): num_codes-dim Q network as output
                      if StackedGAN or 2 Q networks if InfoGAN
  
      Returns:
          Model: Discriminator Model
      """
      kernel_size = 5
      layer_filters = [32, 64, 128, 256]
  
      x = inputs
      for filters in layer_filters:
          # first 3 convolution layers use strides = 2
          # last one uses strides = 1
          if filters == layer_filters[-1]:
              strides = 1
          else:
              strides = 2
          x = LeakyReLU(alpha=0.2)(x)
          x = Conv2D(filters=filters,
                     kernel_size=kernel_size,
                     strides=strides,
                     padding='same')(x)
  
      x = Flatten()(x)
      # default output is probability that the image is real
      outputs = Dense(1)(x)
      if activation is not None:
          print(activation)
          outputs = Activation(activation)(outputs)
  
      if num_labels:
          # ACGAN and InfoGAN have 2nd output
          # 2nd output is 10-dim one-hot vector of label
          layer = Dense(layer_filters[-2])(x)
          labels = Dense(num_labels)(layer)
          labels = Activation('softmax', name='label')(labels)
          if num_codes is None:
              outputs = [outputs, labels]
          else:
              # InfoGAN have 3rd and 4th outputs
              # 3rd output is 1-dim continous Q of 1st c given x
              code1 = Dense(1)(layer)
              code1 = Activation('sigmoid', name='code1')(code1)
  
              # 4th output is 1-dim continuous Q of 2nd c given x
              code2 = Dense(1)(layer)
              code2 = Activation('sigmoid', name='code2')(code2)
  
              outputs = [outputs, labels, code1, code2]
      elif num_codes is not None:
          # StackedGAN Q0 output
          # z0_recon is reconstruction of z0 normal distribution
          z0_recon =  Dense(num_codes)(x)
          z0_recon = Activation('tanh', name='z0')(z0_recon)
          outputs = [outputs, z0_recon]
  
      return Model(inputs, outputs, name='discriminator')
  
  
  def train(models, x_train, params):
      """Train the Discriminator and Adversarial Networks
  
      Alternately train Discriminator and Adversarial networks by batch.
      Discriminator is trained first with properly real and fake images.
      Adversarial is trained next with fake images pretending to be real
      Generate sample images per save_interval.
  
      # Arguments
          models (list): Generator, Discriminator, Adversarial models
          x_train (tensor): Train images
          params (list) : Networks parameters
  
      """
      # the GAN models
      generator, discriminator, adversarial = models
      # network parameters
      batch_size, latent_size, train_steps, model_name = params
      # the generator image is saved every 500 steps
      save_interval = 500
      # noise vector to see how the generator output
      # evolves during training
      noise_input = np.random.uniform(-1.0, 1.0, size=[16, latent_size])
      # number of elements in train dataset
      train_size = x_train.shape[0]
      for i in range(train_steps):
          # train the discriminator for 1 batch
          # 1 batch of real (label=1.0) and fake images (label=0.0)
          # randomly pick real images from dataset
          rand_indexes = np.random.randint(0,
                                           train_size,
                                           size=batch_size)
          real_images = x_train[rand_indexes]
          # generate fake images from noise using generator
          # generate noise using uniform distribution
          noise = np.random.uniform(-1.0,
                                    1.0,
                                    size=[batch_size, latent_size])
          # generate fake images
          fake_images = generator.predict(noise)
          # real + fake images = 1 batch of train data
          x = np.concatenate((real_images, fake_images))
          # label real and fake images
          # real images label is 1.0
          y = np.ones([2 * batch_size, 1])
          # fake images label is 0.0
          y[batch_size:, :] = 0.0
          # train discriminator network, log the loss and accuracy
          loss, acc = discriminator.train_on_batch(x, y)
          log = "%d: [discriminator loss: %f, acc: %f]" % (i, loss, acc)
  
          # train the adversarial network for 1 batch
          # 1 batch of fake images with label=1.0
          # since the discriminator weights are frozen
          # in adversarial network only the generator is trained
          # generate noise using uniform distribution
          noise = np.random.uniform(-1.0,
                                    1.0,
                                    size=[batch_size, latent_size])
          # label fake images as real or 1.0
          y = np.ones([batch_size, 1])
          # train the adversarial network
          # note that unlike in discriminator training,
          # we do not save the fake images in a variable
          # the fake images go to the discriminator input
          # of the adversarial for classification
          # log the loss and accuracy
          loss, acc = adversarial.train_on_batch(noise, y)
          log = "%s [adversarial loss: %f, acc: %f]" % (log, loss, acc)
          if (i + 1) % 25 == 0:
              # plot generator images on a periodic basis
              print(log)
  
      # save the model after training the generator
      # the trained generator can be reloaded
      # for future MNIST digit generation
      generator.save(model_name + ".h5")
  
  
  
  
  def plot_images(generator,
                  noise_input,
                  noise_label=None,
                  noise_codes=None,
                  show=False,
                  step=0,
                  model_name="gan"):
      """Generate fake images and plot them
  
      For visualization purposes, generate fake images
      then plot them in a square grid
  
      # Arguments
          generator (Model): The Generator Model for
              fake images generation
          noise_input (ndarray): Array of z-vectors
          show (bool): Whether to show plot or not
          step (int): Appended to filename of the save images
          model_name (string): Model name
  
      """
      rows = int(math.sqrt(noise_input.shape[0]))
      if noise_label is not None:
          noise_input = [noise_input, noise_label]
          if noise_codes is not None:
              noise_input += noise_codes
  
      images = generator.predict(noise_input)
      plt.figure(figsize=(2.2, 2.2))
      num_images = images.shape[0]
      image_size = images.shape[1]
      for i in range(num_images):
          plt.subplot(rows, rows, i + 1)
          image = np.reshape(images[i], [image_size, image_size])
          plt.imshow(image, cmap='gray')
          plt.axis('off')
      if show:
          plt.show()
      else:
          plt.close('all')
  
  
  def test_generator(generator):
      noise_input = np.random.uniform(-1.0, 1.0, size=[16, 100])
      plot_images(generator,
                  noise_input=noise_input,
                  show=True,
                  model_name="test_outputs")
  
  
  def build_and_train_models():
      """Load the dataset, build LSGAN discriminator,
      generator, and adversarial models.
      Call the LSGAN train routine.
      """
      # load MNIST dataset
      (x_train, _), (_, _) = mnist.load_data()
      # reshape data for CNN as (28, 28, 1) and normalize
      image_size = x_train.shape[1]
      x_train = np.reshape(x_train,
                           [-1, image_size, image_size, 1])
      x_train = x_train.astype('float32') / 255
      model_name = "lsgan_mnist"
      # network parameters
      # the latent or z vector is 100-dim
      latent_size = 100
      input_shape = (image_size, image_size, 1)
      batch_size = 64
      lr = 2e-4
      decay = 6e-8
      train_steps = 40000
      # build discriminator model
      inputs = Input(shape=input_shape, name='discriminator_input')
      discriminator = build_discriminator(inputs, activation=None)
      # [1] uses Adam, but discriminator easily
      # converges with RMSprop
      optimizer = RMSprop(lr=lr, decay=decay)
      # LSGAN uses MSE loss [2]
      discriminator.compile(loss='mse',
                            optimizer=optimizer,
                            metrics=['accuracy'])
      discriminator.summary()
      # build generator model
      input_shape = (latent_size, )
      inputs = Input(shape=input_shape, name='z_input')
      generator = build_generator(inputs, image_size)
      generator.summary()
      # build adversarial model = generator + discriminator
      optimizer = RMSprop(lr=lr*0.5, decay=decay*0.5)
      # freeze the weights of discriminator
      # during adversarial training
      discriminator.trainable = False
      adversarial = Model(inputs,
                          discriminator(generator(inputs)),
                          name=model_name)
      # LSGAN uses MSE loss [2]
      adversarial.compile(loss='mse',
                          optimizer=optimizer,
                          metrics=['accuracy'])
      adversarial.summary()
      # train discriminator and adversarial networks
      models = (generator, discriminator, adversarial)
      params = (batch_size, latent_size, train_steps, model_name)
      train(models, x_train, params)
  
      return models
  
  (g, d, a) = build_and_train_models()
  