r"""Trains a Self-Attention GAN using Estimators.

Code borrowed from tensorlow-gan library.
We do not claim any ownership on this code and you should refer to the LICENCE
of the tensorflow-gan library.

"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl import app
from absl import flags
import train_experiment


flags.DEFINE_string('model_dir', '/tmp/tfgan_logdir/sagan-estimator',
                    'Optional location to save model. If `None`, use a '
                    'default provided by tf.Estimator.')


# ML Hparams.
flags.DEFINE_integer(
    'train_batch_size', 32,
    'The number of images in each train batch. From go/tpu-pitfalls: "The '
    'batch size of any model should always be at least 64 (8 per TPU core), '
    'since the TPU always pads the tensors to this size. The ideal batch size '
    'when training on the TPU is 1024 (128 per TPU core), since this '
    'eliminates inefficiencies related to memory transfer and padding."')
flags.DEFINE_integer('z_dim', 128,
                     'Dimensions of the generator noise vector.')
flags.DEFINE_integer('gf_dim', 64, 'Dimensionality of gf. [64]')
flags.DEFINE_integer('df_dim', 64, 'Dimensionality of df. [64]')
flags.DEFINE_float('generator_lr', 0.0001, 'The generator learning rate.')
flags.DEFINE_float('discriminator_lr', 0.0004,
                   'The discriminator learning rate.')
flags.DEFINE_float('beta1', 0.0, 'Momentum term of adam. [0.0]')
flags.DEFINE_integer('num_classes', 1000, 'Number of classes of the dataset.')


# ML Infra.
flags.DEFINE_enum(
    'mode', None, ['train', 'continuous_eval', 'train_and_eval'],
    'Mode to run in. `train` just trains the model. `continuous_eval` '
    'continuously looks for new checkpoints and computes eval metrics and '
    'writes sample outputs to disk. `train_and_eval` does both. '
    'If not set, will deduce mode from the TF_CONFIG environment variable.')
flags.DEFINE_integer('max_number_of_steps', 50000,
                     'The maximum number of train steps.')
flags.DEFINE_integer(
    'train_steps_per_eval', 1000,
    'Number of train steps before writing some sample images.')
flags.DEFINE_integer('num_eval_steps', 32, 'The number of evaluation steps.')
flags.DEFINE_integer('eval_batch_size', 1024,
                     'The number of images in each eval batch.')
flags.DEFINE_integer('predict_batch_size', 80,
                     'The number of images in each predict batch.')

# Debugging.
flags.DEFINE_bool('use_tpu', True, 'Whether to use TPU or CPU.')
flags.DEFINE_bool('eval_on_tpu', True, 'Whether eval is run on TPU.')
flags.DEFINE_integer(
    'continuous_eval_timeout_secs', None,
    'None, or number of seconds to wait for a checkpoint '
    'before stopping.')

# TPU params.
flags.DEFINE_bool(
    'use_tpu_estimator', True,
    'Whether to use TPUGANEstimator or GANEstimator. This is useful if, for '
    'instance, we want to run the eval job on GPU.')
flags.DEFINE_string('tpu', None, 'A string corresponding to the TPU to use.')
flags.DEFINE_string('gcp_project', None,
                    'Name of the GCP project containing Cloud TPUs.')
flags.DEFINE_string('tpu_zone', None, 'Zone where the TPUs are located.')
flags.DEFINE_integer(
    'tpu_iterations_per_loop', 1000,
    'Steps per interior TPU loop. Should be less than '
    '--train_steps_per_eval.')

FLAGS = flags.FLAGS


def main(_):
    tpu_location = FLAGS.tpu
    hparams = train_experiment.HParams(
        train_batch_size=FLAGS.train_batch_size,
        eval_batch_size=FLAGS.eval_batch_size,
        predict_batch_size=FLAGS.predict_batch_size,
        generator_lr=FLAGS.generator_lr,
        discriminator_lr=FLAGS.discriminator_lr,
        beta1=FLAGS.beta1,
        gf_dim=FLAGS.gf_dim,
        df_dim=FLAGS.df_dim,
        num_classes=FLAGS.num_classes,
        shuffle_buffer_size=10000,
        z_dim=FLAGS.z_dim,
        model_dir=FLAGS.model_dir,
        max_number_of_steps=FLAGS.max_number_of_steps,
        train_steps_per_eval=FLAGS.train_steps_per_eval,
        num_eval_steps=FLAGS.num_eval_steps,
        debug_params=train_experiment.DebugParams(
            use_tpu=FLAGS.use_tpu,
            eval_on_tpu=FLAGS.eval_on_tpu,
            fake_nets=False,
            fake_data=False,
            continuous_eval_timeout_secs=FLAGS.continuous_eval_timeout_secs,
        ),
        tpu_params=train_experiment.TPUParams(
            use_tpu_estimator=FLAGS.use_tpu_estimator,
            tpu_location=tpu_location,
            gcp_project=FLAGS.gcp_project,
            tpu_zone=FLAGS.tpu_zone,
            tpu_iterations_per_loop=FLAGS.tpu_iterations_per_loop,
        ),
    )
    if FLAGS.mode == 'train':
        train_experiment.run_train(hparams)
    elif FLAGS.mode == 'continuous_eval':
        train_experiment.run_continuous_eval(hparams)
    elif FLAGS.mode == 'train_and_eval' or FLAGS.mode is None:
        train_experiment.run_train_and_eval(hparams)
    else:
        raise ValueError('Mode not recognized: ', FLAGS.mode)


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