import os
import tensorflow as tf
from util import constants
from util.config_util import get_model_params, get_task_params, get_train_params
from tf2_models.trainer import Trainer
from absl import app
from absl import flags

from util.models import MODELS
from util.tasks import TASKS

FLAGS = flags.FLAGS
flags.DEFINE_string('exp_name', 'trial1', 'experiment directory')
flags.DEFINE_string('task', 'word_sv_agreement_lm', 'sv_agreement_lm | word_sv_agreement_lm')
flags.DEFINE_string('model', 'lm_lstm', 'lm_lstm | lm_gpt2 | lm_gpt2_shared | lm_lstm_shared_emb | cl_gpt2_shared | cl_gpt2 | cl_lstm')
flags.DEFINE_string('model_config', 'base', 'base | small_lstm ')
flags.DEFINE_string('train_config', 'radam_fast', 'radam_slow | radam_fast')
flags.DEFINE_integer('keep_checkpoint_every_n_hours',None, 'keep_checkpoint_every_n_hours passed to training manager')
flags.DEFINE_integer('batch_size',16, 'batch_size')


hparams = flags.FLAGS


def run():
  gpus = tf.config.experimental.list_physical_devices('GPU')
  if gpus:
    # Currently, memory growth needs to be the same across GPUs
    try:
      for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
      print(e)

  strategy = tf.distribute.MirroredStrategy()



  log_dir = "logs"
  chkpt_dir = "tf_ckpts"

  # Create task
  with strategy.scope():
    task = TASKS[hparams.task](get_task_params())

    # Create the Model
    model_params = get_model_params(task,hparams.model, hparams.model_config)
    print("model_params: ", model_params.__dict__)

    model = MODELS[hparams.model](hparams=get_model_params(task,hparams.model, hparams.model_config))

  trainer_params = get_train_params(hparams.train_config)

  log_dir = os.path.join(log_dir,task.name, model.model_name+"_"+str(hparams.model_config)+"_"+str(trainer_params.learning_rate)+"_"+hparams.exp_name)
  ckpt_dir = os.path.join(chkpt_dir,task.name, model.model_name+"_"+str(hparams.model_config)+"_"+str(trainer_params.learning_rate)+"_"+hparams.exp_name)

  # Create task
  trainer = Trainer(hparams,
                    strategy=strategy,
                    task=task,
                    model=model,
                    train_params=trainer_params,
                    log_dir=log_dir,
                    ckpt_dir=ckpt_dir)

  trainer.restore()
  trainer.train()


def main(argv):
  if len(argv) > 1:
    raise app.UsageError('Too many command-line arguments.')

  run()

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