text: from __future__ import division
  from __future__ import print_function
  
  import prettytensor as pt
  import tensorflow as tf
  import numpy as np
  import scipy.misc
  import os
  import sys
  from six.moves import range
  from progressbar import ETA, Bar, Percentage, ProgressBar
  
  
  from misc.config import cfg
  from misc.utils import mkdir_p
  
  TINY = 1e-8
  
  
  # reduce_mean normalize also the dimension of the embeddings
  def KL_loss(mu, log_sigma):
      with tf.name_scope("KL_divergence"):
          loss = -log_sigma + .5 * (-1 + tf.exp(2. * log_sigma) + tf.square(mu))
          loss = tf.reduce_mean(loss)
          return loss
  
  
  class CondGANTrainer(object):
      def __init__(self,
                   model,
                   models_dir,
                   dataset=None,
                   exp_name="model",
                   ckt_logs_dir="ckt_logs",
                   ):
          """
          :type model: RegularizedGAN
          """
          self.model = model
          self.dataset = dataset
          self.exp_name = exp_name
          self.log_dir = ckt_logs_dir
          self.checkpoint_dir = models_dir
  
          self.batch_size = cfg.TRAIN.BATCH_SIZE
          self.max_epoch = cfg.TRAIN.MAX_EPOCH
          self.snapshot_interval = cfg.TRAIN.SNAPSHOT_INTERVAL
          self.model_path = cfg.TRAIN.PRETRAINED_MODEL
  
          self.log_vars = []
  
      def build_placeholder(self):
          '''Helper function for init_opt'''
          self.images = tf.placeholder(
              tf.float32, [self.batch_size] + self.dataset.image_shape,
              name='real_images')
          self.wrong_images = tf.placeholder(
              tf.float32, [self.batch_size] + self.dataset.image_shape,
              name='wrong_images'
          )
          self.embeddings = tf.placeholder(
              tf.float32, [self.batch_size] + self.dataset.embedding_shape,
              name='conditional_embeddings'
          )
  
          self.generator_lr = tf.placeholder(
              tf.float32, [],
              name='generator_learning_rate'
          )
          self.discriminator_lr = tf.placeholder(
              tf.float32, [],
              name='discriminator_learning_rate'
          )
  
      def sample_encoded_context(self, embeddings):
          '''Helper function for init_opt'''
          c_mean_logsigma = self.model.generate_condition(embeddings)
          mean = c_mean_logsigma[0]
          if cfg.TRAIN.COND_AUGMENTATION:
              # epsilon = tf.random_normal(tf.shape(mean))
              epsilon = tf.truncated_normal(tf.shape(mean))
              stddev = tf.exp(c_mean_logsigma[1])
              c = mean + stddev * epsilon
  
              kl_loss = KL_loss(c_mean_logsigma[0], c_mean_logsigma[1])
          else:
              c = mean
              kl_loss = 0
  
          return c, cfg.TRAIN.COEFF.KL * kl_loss
  
      def init_opt(self):
          self.build_placeholder()
  
          with pt.defaults_scope(phase=pt.Phase.train):
              with tf.variable_scope("g_net"):
                  # ####get output from G network################################
                  c, kl_loss = self.sample_encoded_context(self.embeddings)
                  z = tf.random_normal([self.batch_size, cfg.Z_DIM])
                  self.log_vars.append(("hist_c", c))
                  self.log_vars.append(("hist_z", z))
                  fake_images = self.model.get_generator(tf.concat(1, [c, z]))
  
              # ####get discriminator_loss and generator_loss ###################
              discriminator_loss, generator_loss =\
                  self.compute_losses(self.images,
                                      self.wrong_images,
                                      fake_images,
                                      self.embeddings)
              generator_loss += kl_loss
              self.log_vars.append(("g_loss_kl_loss", kl_loss))
              self.log_vars.append(("g_loss", generator_loss))
              self.log_vars.append(("d_loss", discriminator_loss))
  
              # #######Total loss for build optimizers###########################
              self.prepare_trainer(generator_loss, discriminator_loss)
              # #######define self.g_sum, self.d_sum,....########################
              self.define_summaries()
  
          with pt.defaults_scope(phase=pt.Phase.test):
              with tf.variable_scope("g_net", reuse=True):
                  self.sampler()
              self.visualization(cfg.TRAIN.NUM_COPY)
              print("success")
  
      def sampler(self):
          c, _ = self.sample_encoded_context(self.embeddings)
          if cfg.TRAIN.FLAG:
              z = tf.zeros([self.batch_size, cfg.Z_DIM])  # Expect similar BGs
          else:
              z = tf.random_normal([self.batch_size, cfg.Z_DIM])
          self.fake_images = self.model.get_generator(tf.concat(1, [c, z]))
  
      def compute_losses(self, images, wrong_images, fake_images, embeddings):
          real_logit = self.model.get_discriminator(images, embeddings)
          wrong_logit = self.model.get_discriminator(wrong_images, embeddings)
          fake_logit = self.model.get_discriminator(fake_images, embeddings)
  
          real_d_loss =\
              tf.nn.sigmoid_cross_entropy_with_logits(real_logit,
                                                      tf.ones_like(real_logit))
          real_d_loss = tf.reduce_mean(real_d_loss)
          wrong_d_loss =\
              tf.nn.sigmoid_cross_entropy_with_logits(wrong_logit,
                                                      tf.zeros_like(wrong_logit))
          wrong_d_loss = tf.reduce_mean(wrong_d_loss)
          fake_d_loss =\
              tf.nn.sigmoid_cross_entropy_with_logits(fake_logit,
                                                      tf.zeros_like(fake_logit))
          fake_d_loss = tf.reduce_mean(fake_d_loss)
          if cfg.TRAIN.B_WRONG:
              discriminator_loss =\
                  real_d_loss + (wrong_d_loss + fake_d_loss) / 2.
              self.log_vars.append(("d_loss_wrong", wrong_d_loss))
          else:
              discriminator_loss = real_d_loss + fake_d_loss
          self.log_vars.append(("d_loss_real", real_d_loss))
          self.log_vars.append(("d_loss_fake", fake_d_loss))
  
          generator_loss = \
              tf.nn.sigmoid_cross_entropy_with_logits(fake_logit,
                                                      tf.ones_like(fake_logit))
          generator_loss = tf.reduce_mean(generator_loss)
  
          return discriminator_loss, generator_loss
  
      def prepare_trainer(self, generator_loss, discriminator_loss):
          '''Helper function for init_opt'''
          all_vars = tf.trainable_variables()
  
          g_vars = [var for var in all_vars if
                    var.name.startswith('g_')]
          d_vars = [var for var in all_vars if
                    var.name.startswith('d_')]
  
          generator_opt = tf.train.AdamOptimizer(self.generator_lr,
                                                 beta1=0.5)
          self.generator_trainer =\
              pt.apply_optimizer(generator_opt,
                                 losses=[generator_loss],
                                 var_list=g_vars)
          discriminator_opt = tf.train.AdamOptimizer(self.discriminator_lr,
                                                     beta1=0.5)
          self.discriminator_trainer =\
              pt.apply_optimizer(discriminator_opt,
                                 losses=[discriminator_loss],
                                 var_list=d_vars)
          self.log_vars.append(("g_learning_rate", self.generator_lr))
          self.log_vars.append(("d_learning_rate", self.discriminator_lr))
  
      def define_summaries(self):
          '''Helper function for init_opt'''
          all_sum = {'g': [], 'd': [], 'hist': []}
          for k, v in self.log_vars:
              if k.startswith('g'):
                  all_sum['g'].append(tf.scalar_summary(k, v))
              elif k.startswith('d'):
                  all_sum['d'].append(tf.scalar_summary(k, v))
              elif k.startswith('hist'):
                  all_sum['hist'].append(tf.histogram_summary(k, v))
  
          self.g_sum = tf.merge_summary(all_sum['g'])
          self.d_sum = tf.merge_summary(all_sum['d'])
          self.hist_sum = tf.merge_summary(all_sum['hist'])
  
      def visualize_one_superimage(self, img_var, images, rows, filename):
          stacked_img = []
          for row in range(rows):
              img = images[row * rows, :, :, :]
              row_img = [img]  # real image
              for col in range(rows):
                  row_img.append(img_var[row * rows + col, :, :, :])
              # each rows is 1realimage +10_fakeimage
              stacked_img.append(tf.concat(1, row_img))
          imgs = tf.expand_dims(tf.concat(0, stacked_img), 0)
          current_img_summary = tf.image_summary(filename, imgs)
          return current_img_summary, imgs
  
      def visualization(self, n):
          fake_sum_train, superimage_train = \
              self.visualize_one_superimage(self.fake_images[:n * n],
                                            self.images[:n * n],
                                            n, "train")
          fake_sum_test, superimage_test = \
              self.visualize_one_superimage(self.fake_images[n * n:2 * n * n],
                                            self.images[n * n:2 * n * n],
                                            n, "test")
          self.superimages = tf.concat(0, [superimage_train, superimage_test])
          self.image_summary = tf.merge_summary([fake_sum_train, fake_sum_test])
  
      def preprocess(self, x, n):
          # make sure every row with n column have the same embeddings
          for i in range(n):
              for j in range(1, n):
                  x[i * n + j] = x[i * n]
          return x
  
      def epoch_sum_images(self, sess, n):
          images_train, _, embeddings_train, captions_train, _ =\
              self.dataset.train.next_batch(n * n, cfg.TRAIN.NUM_EMBEDDING)
          images_train = self.preprocess(images_train, n)
          embeddings_train = self.preprocess(embeddings_train, n)
  
          images_test, _, embeddings_test, captions_test, _ = \
              self.dataset.test.next_batch(n * n, 1)
          images_test = self.preprocess(images_test, n)
          embeddings_test = self.preprocess(embeddings_test, n)
  
          images = np.concatenate([images_train, images_test], axis=0)
          embeddings =\
              np.concatenate([embeddings_train, embeddings_test], axis=0)
  
          if self.batch_size > 2 * n * n:
              images_pad, _, embeddings_pad, _, _ =\
                  self.dataset.test.next_batch(self.batch_size - 2 * n * n, 1)
              images = np.concatenate([images, images_pad], axis=0)
              embeddings = np.concatenate([embeddings, embeddings_pad], axis=0)
          feed_dict = {self.images: images,
                       self.embeddings: embeddings}
          gen_samples, img_summary =\
              sess.run([self.superimages, self.image_summary], feed_dict)
  
          # save images generated for train and test captions
          scipy.misc.imsave('%s/train.jpg' % (self.log_dir), gen_samples[0])
          scipy.misc.imsave('%s/test.jpg' % (self.log_dir), gen_samples[1])
  
          # pfi_train = open(self.log_dir + "/train.txt", "w")
          pfi_test = open(self.log_dir + "/test.txt", "w")
          for row in range(n):
              # pfi_train.write('\n***row %d***\n' % row)
              # pfi_train.write(captions_train[row * n])
  
              pfi_test.write('\n***row %d***\n' % row)
              pfi_test.write(captions_test[row * n])
          # pfi_train.close()
          pfi_test.close()
  
          return img_summary
  
      def build_model(self, sess):
          self.init_opt()
          sess.run(tf.initialize_all_variables())
  
          if len(self.model_path) > 0:
              print("Reading model parameters from %s" % self.model_path)
              restore_vars = tf.all_variables()
              # all_vars = tf.all_variables()
              # restore_vars = [var for var in all_vars if
              #                 var.name.startswith('g_') or
              #                 var.name.startswith('d_')]
              saver = tf.train.Saver(restore_vars)
              saver.restore(sess, self.model_path)
  
              istart = self.model_path.rfind('_') + 1
              iend = self.model_path.rfind('.')
              counter = self.model_path[istart:iend]
              counter = int(counter)
          else:
              print("Created model with fresh parameters.")
              counter = 0
          return counter
  
      def train(self):
          config = tf.ConfigProto(allow_soft_placement=True)
          with tf.Session(config=config) as sess:
              with tf.device("/gpu:%d" % cfg.GPU_ID):
                  counter = self.build_model(sess)
                  saver = tf.train.Saver(tf.all_variables(),
                                         keep_checkpoint_every_n_hours=2)
  
                  # summary_op = tf.merge_all_summaries()
                  summary_writer = tf.train.SummaryWriter(self.log_dir,
                                                          sess.graph)
  
                  keys = ["d_loss", "g_loss"]
                  log_vars = []
                  log_keys = []
                  for k, v in self.log_vars:
                      if k in keys:
                          log_vars.append(v)
                          log_keys.append(k)
                          # print(k, v)
                  generator_lr = cfg.TRAIN.GENERATOR_LR
                  discriminator_lr = cfg.TRAIN.DISCRIMINATOR_LR
                  num_embedding = cfg.TRAIN.NUM_EMBEDDING
                  lr_decay_step = cfg.TRAIN.LR_DECAY_EPOCH
                  number_example = self.dataset.train._num_examples
                  updates_per_epoch = int(number_example / self.batch_size)
                  epoch_start = int(counter / updates_per_epoch)
                  for epoch in range(epoch_start, self.max_epoch):
                      widgets = ["epoch #%d|" % epoch,
                                 Percentage(), Bar(), ETA()]
                      pbar = ProgressBar(maxval=updates_per_epoch,
                                         widgets=widgets)
                      pbar.start()
  
                      if epoch % lr_decay_step == 0 and epoch != 0:
                          generator_lr *= 0.5
                          discriminator_lr *= 0.5
  
                      all_log_vals = []
                      for i in range(updates_per_epoch):
                          pbar.update(i)
                          # training d
                          images, wrong_images, embeddings, _, _ =\
                              self.dataset.train.next_batch(self.batch_size,
                                                            num_embedding)
                          feed_dict = {self.images: images,
                                       self.wrong_images: wrong_images,
                                       self.embeddings: embeddings,
                                       self.generator_lr: generator_lr,
                                       self.discriminator_lr: discriminator_lr
                                       }
                          # train d
                          feed_out = [self.discriminator_trainer,
                                      self.d_sum,
                                      self.hist_sum,
                                      log_vars]
                          _, d_sum, hist_sum, log_vals = sess.run(feed_out,
                                                                  feed_dict)
                          summary_writer.add_summary(d_sum, counter)
                          summary_writer.add_summary(hist_sum, counter)
                          all_log_vals.append(log_vals)
                          # train g
                          feed_out = [self.generator_trainer,
                                      self.g_sum]
                          _, g_sum = sess.run(feed_out,
                                              feed_dict)
                          summary_writer.add_summary(g_sum, counter)
                          # save checkpoint
                          counter += 1
                          if counter % self.snapshot_interval == 0:
                              snapshot_path = "%s/%s.ckpt" %\
                                               (self.checkpoint_dir,
                                                self.exp_name)
                              fn = saver.save(sess, snapshot_path)
                              print("Model saved in file: %s" % fn)
  
                      img_sum = self.epoch_sum_images(sess, cfg.TRAIN.NUM_COPY)
                      summary_writer.add_summary(img_sum, counter)
  
                      avg_log_vals = np.mean(np.array(all_log_vals), axis=0)
                      dic_logs = {}
                      for k, v in zip(log_keys, avg_log_vals):
                          dic_logs[k] = v
                          # print(k, v)
  
                      log_line = "; ".join("%s: %s" %
                                           (str(k), str(dic_logs[k]))
                                           for k in dic_logs)
                      print("Epoch %d | " % (epoch) + log_line)
                      sys.stdout.flush()
                      if np.any(np.isnan(avg_log_vals)):
                          raise ValueError("NaN detected!")
  
      def save_super_images(self, images, sample_batchs, filenames,
                            sentenceID, save_dir, subset):
          # batch_size samples for each embedding
          numSamples = len(sample_batchs)
          for j in range(len(filenames)):
              s_tmp = '%s-1real-%dsamples/%s/%s' %\
                  (save_dir, numSamples, subset, filenames[j])
              folder = s_tmp[:s_tmp.rfind('/')]
              if not os.path.isdir(folder):
                  print('Make a new folder: ', folder)
                  mkdir_p(folder)
              superimage = [images[j]]
              # cfg.TRAIN.NUM_COPY samples for each text embedding/sentence
              for i in range(len(sample_batchs)):
                  superimage.append(sample_batchs[i][j])
  
              superimage = np.concatenate(superimage, axis=1)
              fullpath = '%s_sentence%d.jpg' % (s_tmp, sentenceID)
              scipy.misc.imsave(fullpath, superimage)
  
      def eval_one_dataset(self, sess, dataset, save_dir, subset='train'):
          count = 0
          print('num_examples:', dataset._num_examples)
          while count < dataset._num_examples:
              start = count % dataset._num_examples
              images, embeddings_batchs, filenames, _ =\
                  dataset.next_batch_test(self.batch_size, start, 1)
              print('count = ', count, 'start = ', start)
              for i in range(len(embeddings_batchs)):
                  samples_batchs = []
                  # Generate up to 16 images for each sentence,
                  # with randomness from noise z and conditioning augmentation.
                  for j in range(np.minimum(16, cfg.TRAIN.NUM_COPY)):
                      samples = sess.run(self.fake_images,
                                         {self.embeddings: embeddings_batchs[i]})
                      samples_batchs.append(samples)
                  self.save_super_images(images, samples_batchs,
                                         filenames, i, save_dir,
                                         subset)
  
              count += self.batch_size
  
      def evaluate(self):
          config = tf.ConfigProto(allow_soft_placement=True)
          with tf.Session(config=config) as sess:
              with tf.device("/gpu:%d" % cfg.GPU_ID):
                  if self.model_path.find('.ckpt') != -1:
                      self.init_opt()
                      print("Reading model parameters from %s" % self.model_path)
                      saver = tf.train.Saver(tf.all_variables())
                      saver.restore(sess, self.model_path)
                      # self.eval_one_dataset(sess, self.dataset.train,
                      #                       self.log_dir, subset='train')
                      self.eval_one_dataset(sess, self.dataset.test,
                                            self.log_dir, subset='test')
                  else:
                      print("Input a valid model path.")
  