"""
Andrin Jenal, 2017
ETH Zurich
"""

import tensorflow as tf

from vae_dcgan import VAE_DCGAN
import hdf5_dataset
from checkpoint_saver import CheckpointSaver
from visualizer import ImageVisualizer


flags = tf.app.flags
flags.DEFINE_string("dataset", "datasets/celeb_dataset_3k_colored.h5", "sample results dir")
flags.DEFINE_boolean("binarized", False, "data set binarization")
flags.DEFINE_string("data_dir", "results/", "checkpoint and logging results dir")
flags.DEFINE_integer("batch_size", 128, "batch size")
flags.DEFINE_integer("image_size", 64, "image size")
flags.DEFINE_integer("channels", 3, "color channels")
flags.DEFINE_integer("max_epoch", 500, "max epoch")
flags.DEFINE_integer("z_size", 100, "size of latent (feature?) space")
flags.DEFINE_float("learning_rate_enc", 5e-4, "learning rate")
flags.DEFINE_float("learning_rate_dec", 5e-4, "learning rate")
flags.DEFINE_float("learning_rate_dis", 5e-4, "learning rate")
flags.DEFINE_integer("generation_step", 1, "generate random images")
flags.DEFINE_integer("checkpoint_step", 10, "save checkpoints")
FLAGS = flags.FLAGS


def main(_):
    # create checkpoint saver
    # the checkpoint saver, can create checkpoint files, which later can be use to restore a model state, but it also
    # audits the model progress to a log file
    checkpoint_saver = CheckpointSaver(FLAGS.data_dir, experiment_name="VAE_DCGAN")
    checkpoint_saver.save_experiment_config(FLAGS.__dict__['__flags'])

    # load training data
    data_set, data_set_shape = hdf5_dataset.read_data_set(FLAGS.dataset, image_size=FLAGS.image_size, shape=(FLAGS.image_size, FLAGS.image_size, FLAGS.channels), binarized=FLAGS.binarized, validation=0)
    train_data = data_set.train

    # create a data visualizer
    visualizer = ImageVisualizer(checkpoint_saver.get_experiment_dir(), image_size=FLAGS.image_size)
    visualizer.training_data_sample(train_data)

    # create the actual DCGAN model
    dcgan_model = VAE_DCGAN(FLAGS.image_size, FLAGS.channels, z_size=FLAGS.z_size, learning_rate_enc=FLAGS.learning_rate_enc, learning_rate_dis=FLAGS.learning_rate_dis, learning_rate_dec=FLAGS.learning_rate_dec)

    print("start", type(dcgan_model).__name__, "model training")
    with tf.Session() as sess:

        sess.run(tf.global_variables_initializer())
        dcgan_model.initialize_summaries(sess, checkpoint_saver.get_experiment_dir())

        for epoch in range(FLAGS.max_epoch):

            for images in train_data.next_batch(FLAGS.batch_size):
                kl, d_loss, g_loss, lth_loss = dcgan_model.update_params(sess, images)

            msg = "epoch: %3d" % epoch + " kl loss  %.4f" % kl + " discriminator loss %.4f" % d_loss + " generator loss  %.4f" % g_loss + " lth layer loss  %.4f" % lth_loss
            checkpoint_saver.audit_loss(msg)

            dcgan_model.update_summaries(sess, images, epoch)

            if epoch % FLAGS.generation_step == 0:
                visualizer.save_generated_samples(dcgan_model.generate_samples(sess, num_samples=200), epoch)

            if epoch % FLAGS.checkpoint_step == 0:
                checkpoint_saver.save_checkpoint(dcgan_model.saver, sess, epoch)

if __name__ == '__main__':
    tf.app.run()

@@ -838,19 +838,19 @@
 ng_rate_
-dec
+gen
 %22, 5e-4,
@@ -1039,17 +1039,17 @@
 _step%22, 
-1
+5
 0, %22save
@@ -2133,19 +2133,19 @@
 ng_rate_
-dec
+gen
 )%0A%0A    p
@@ -2605,16 +2605,17 @@
 och: %253d
+,
 %22 %25 epoc
@@ -2624,25 +2624,24 @@
 + %22 kl loss 
- 
 %25.4f%22 %25 kl +
@@ -2642,16 +2642,17 @@
 %25 kl + %22
+,
  discrim
@@ -2681,16 +2681,17 @@
 loss + %22
+,
  generat
@@ -2694,25 +2694,24 @@
 erator loss 
- 
 %25.4f%22 %25 g_lo
@@ -2716,16 +2716,17 @@
 loss + %22
+,
  lth lay
@@ -2733,17 +2733,16 @@
 er loss 
- 
 %25.4f%22 %25 
@@ -2794,16 +2794,63 @@
 oss(msg)
+%0A            checkpoint_saver.audit_time(epoch)
 %0A%0A      

change defaults, audit time