"""Detection model trainer.

This file provides a generic training method to train a
DetectionModel.
"""
import datetime
import os
import tensorflow as tf
import time

from avod.builders import optimizer_builder
from avod.core import trainer_utils
from avod.core import summary_utils

slim = tf.contrib.slim

def train(model, train_config):
    """Training function for detection models.

    Args:
        model: The detection model object.
        train_config: a train_*pb2 protobuf.
            training i.e. loading RPN weights onto AVOD model.
    """

    model = model
    train_config = train_config
    # Get model configurations
    model_config = model.model_config

    # Create a variable tensor to hold the global step
    global_step_tensor = tf.Variable(
        0, trainable=False, name='global_step')

    #############################
    # Get training configurations
    #############################
    max_iterations = train_config.max_iterations
    summary_interval = train_config.summary_interval
    checkpoint_interval = train_config.checkpoint_interval
    max_checkpoints = train_config.max_checkpoints_to_keep

    paths_config = model_config.paths_config
    logdir = paths_config.logdir
    if not os.path.exists(logdir):
        os.makedirs(logdir)

    checkpoint_dir = paths_config.checkpoint_dir
    if not os.path.exists(checkpoint_dir):
        os.makedirs(checkpoint_dir)
    checkpoint_path = checkpoint_dir + '/' + \
        model_config.checkpoint_name

    pretrained_checkpoint_dir = checkpoint_dir + '/../../' + \
        'pyramid_cars_with_aug_dt_5_tracking_corr_pretrained/checkpoints'

    global_summaries = set([])

    # The model should return a dictionary of predictions
    prediction_dict = model.build()

    summary_histograms = train_config.summary_histograms
    summary_img_images = train_config.summary_img_images
    summary_bev_images = train_config.summary_bev_images

    # get variables to train
    if not train_config.use_pretrained_model:
        variable_to_train = None
    else:
        trainable_variables = tf.trainable_variables()
        variable_to_train = trainable_variables[68:72] + \
                            trainable_variables[96:]


    ##############################
    # Setup loss
    ##############################
    losses_dict, total_loss = model.loss(prediction_dict)

    # Optimizer
    training_optimizer = optimizer_builder.build(
        train_config.optimizer,
        global_summaries,
        global_step_tensor)

    # Create the train op
    with tf.variable_scope('train_op'):
        train_op = slim.learning.create_train_op(
            total_loss,
            training_optimizer,
            variables_to_train=variable_to_train,
            clip_gradient_norm=1.0,
            global_step=global_step_tensor)

    # Add the result of the train_op to the summary
    tf.summary.scalar("training_loss", train_op)

    # Add maximum memory usage summary op
    # This op can only be run on device with gpu
    # so it's skipped on travis
    is_travis = 'TRAVIS' in os.environ
    if not is_travis:
        # tf.summary.scalar('bytes_in_use',
        #                   tf.contrib.memory_stats.BytesInUse())
        tf.summary.scalar('max_bytes',
                          tf.contrib.memory_stats.MaxBytesInUse())

    summaries = set(tf.get_collection(tf.GraphKeys.SUMMARIES))
    summary_merged = summary_utils.summaries_to_keep(
        summaries,
        global_summaries,
        histograms=summary_histograms,
        input_imgs=summary_img_images,
        input_bevs=summary_bev_images
    )

    allow_gpu_mem_growth = train_config.allow_gpu_mem_growth
    if allow_gpu_mem_growth:
        # GPU memory config
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = allow_gpu_mem_growth
        sess = tf.Session(config=config)
    else:
        sess = tf.Session()

    # Create unique folder name using datetime for summary writer
    datetime_str = str(datetime.datetime.now())
    logdir = logdir + '/train'
    train_writer = tf.summary.FileWriter(logdir + '/' + datetime_str,
                                         sess.graph)
    # Save checkpoints regularly.
    saver = tf.train.Saver(max_to_keep=max_checkpoints, pad_step_number=True)

    # Create init op
    # if train_config.use_pretrained_model:
    #     init = tf.initialize_variables(variable_to_train)
    # else:
    #     init = tf.global_variables_initializer()

    init = tf.global_variables_initializer()

    # Continue from last saved checkpoint
    if not train_config.overwrite_checkpoints:
        trainer_utils.load_checkpoints(checkpoint_dir,saver)
        if len(saver.last_checkpoints) > 0:
            checkpoint_to_restore = saver.last_checkpoints[-1]
            saver.restore(sess, checkpoint_to_restore)
        else:
            sess.run(init)
            # load pretrained model
            if train_config.use_pretrained_model:
                variable_to_restore = tf.trainable_variables()
                variable_to_restore = variable_to_restore[:68] + \
                                      variable_to_restore[72:96]
                variable_to_restore = {var.op.name: var for var in variable_to_restore}
                saver2 = tf.train.Saver(var_list=variable_to_restore)
                print('Loading pretrained model...')
                trainer_utils.load_checkpoints(pretrained_checkpoint_dir, saver2)
                checkpoint_to_restore = saver2.last_checkpoints[11]
                saver2.restore(sess, checkpoint_to_restore)
    else:
        sess.run(init)
        # load pretrained model
        if train_config.use_pretrained_model:
            variable_to_restore = tf.trainable_variables()
            variable_to_restore = variable_to_restore[:68] + \
                                  variable_to_restore[72:96]
            variable_to_restore = {var.op.name: var for var in variable_to_restore}
            saver2 = tf.train.Saver(var_list=variable_to_restore)
            print('Loading pretrained model...')
            trainer_utils.load_checkpoints(pretrained_checkpoint_dir, saver2)
            checkpoint_to_restore = saver2.last_checkpoints[11]
            saver2.restore(sess, checkpoint_to_restore)


    # Read the global step if restored
    global_step = tf.train.global_step(sess, global_step_tensor)
    print('Starting from step {} / {}'.format(
        global_step, max_iterations))

    # Main Training Loop
    last_time = time.time()
    for step in range(global_step, max_iterations + 1):

        # Save checkpoint
        if step % checkpoint_interval == 0:
            global_step = tf.train.global_step(sess,
                                               global_step_tensor)

            saver.save(sess,
                       save_path=checkpoint_path,
                       global_step=global_step)

            print('Step {} / {}, Checkpoint saved to {}-{:08d}'.format(
                step, max_iterations,
                checkpoint_path, global_step))

        feed_dict = model.create_feed_dict()

        # Write summaries and train op
        if step % summary_interval == 0:
            current_time = time.time()
            time_elapsed = current_time - last_time
            last_time = current_time

            train_op_loss, summary_out = sess.run(
                    [train_op, summary_merged], feed_dict=feed_dict)
            print('Step {}, Total Loss {:0.3f}, Time Elapsed {:0.3f} s'.format(
                    step, train_op_loss, time_elapsed))

            train_writer.add_summary(summary_out, step)

        else:
            # Run the train op only
            sess.run(train_op, feed_dict)

    # Close the summary writers
    train_writer.close()
[(32, 'tensorflow.Variable', 'tf.Variable', (['(0)'], {'trainable': '(False)', 'name': '"""global_step"""'}), True, 'import tensorflow as tf\n'), (81, 'avod.builders.optimizer_builder.build', 'optimizer_builder.build', (['train_config.optimizer', 'global_summaries', 'global_step_tensor'], {}), False, 'from avod.builders import optimizer_builder\n'), (96, 'tensorflow.summary.scalar', 'tf.summary.scalar', (['"""training_loss"""', 'train_op'], {}), True, 'import tensorflow as tf\n'), (109, 'avod.core.summary_utils.summaries_to_keep', 'summary_utils.summaries_to_keep', (['summaries', 'global_summaries'], {'histograms': 'summary_histograms', 'input_imgs': 'summary_img_images', 'input_bevs': 'summary_bev_images'}), False, 'from avod.core import summary_utils\n'), (129, 'tensorflow.summary.FileWriter', 'tf.summary.FileWriter', (["(logdir + '/' + datetime_str)", 'sess.graph'], {}), True, 'import tensorflow as tf\n'), (132, 'tensorflow.train.Saver', 'tf.train.Saver', ([], {'max_to_keep': 'max_checkpoints', 'pad_step_number': '(True)'}), True, 'import tensorflow as tf\n'), (140, 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', ([], {}), True, 'import tensorflow as tf\n'), (177, 'tensorflow.train.global_step', 'tf.train.global_step', (['sess', 'global_step_tensor'], {}), True, 'import tensorflow as tf\n'), (182, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (45, 'os.path.exists', 'os.path.exists', (['logdir'], {}), False, 'import os\n'), (46, 'os.makedirs', 'os.makedirs', (['logdir'], {}), False, 'import os\n'), (49, 'os.path.exists', 'os.path.exists', (['checkpoint_dir'], {}), False, 'import os\n'), (50, 'os.makedirs', 'os.makedirs', (['checkpoint_dir'], {}), False, 'import os\n'), (70, 'tensorflow.trainable_variables', 'tf.trainable_variables', ([], {}), True, 'import tensorflow as tf\n'), (87, 'tensorflow.variable_scope', 'tf.variable_scope', (['"""train_op"""'], {}), True, 'import tensorflow as tf\n'), (108, 'tensorflow.get_collection', 'tf.get_collection', (['tf.GraphKeys.SUMMARIES'], {}), True, 'import tensorflow as tf\n'), (120, 'tensorflow.ConfigProto', 'tf.ConfigProto', ([], {}), True, 'import tensorflow as tf\n'), (122, 'tensorflow.Session', 'tf.Session', ([], {'config': 'config'}), True, 'import tensorflow as tf\n'), (124, 'tensorflow.Session', 'tf.Session', ([], {}), True, 'import tensorflow as tf\n'), (127, 'datetime.datetime.now', 'datetime.datetime.now', ([], {}), False, 'import datetime\n'), (144, 'avod.core.trainer_utils.load_checkpoints', 'trainer_utils.load_checkpoints', (['checkpoint_dir', 'saver'], {}), False, 'from avod.core import trainer_utils\n'), (106, 'tensorflow.contrib.memory_stats.MaxBytesInUse', 'tf.contrib.memory_stats.MaxBytesInUse', ([], {}), True, 'import tensorflow as tf\n'), (165, 'tensorflow.trainable_variables', 'tf.trainable_variables', ([], {}), True, 'import tensorflow as tf\n'), (169, 'tensorflow.train.Saver', 'tf.train.Saver', ([], {'var_list': 'variable_to_restore'}), True, 'import tensorflow as tf\n'), (171, 'avod.core.trainer_utils.load_checkpoints', 'trainer_utils.load_checkpoints', (['pretrained_checkpoint_dir', 'saver2'], {}), False, 'from avod.core import trainer_utils\n'), (187, 'tensorflow.train.global_step', 'tf.train.global_step', (['sess', 'global_step_tensor'], {}), True, 'import tensorflow as tf\n'), (202, 'time.time', 'time.time', ([], {}), False, 'import time\n'), (152, 'tensorflow.trainable_variables', 'tf.trainable_variables', ([], {}), True, 'import tensorflow as tf\n'), (156, 'tensorflow.train.Saver', 'tf.train.Saver', ([], {'var_list': 'variable_to_restore'}), True, 'import tensorflow as tf\n'), (158, 'avod.core.trainer_utils.load_checkpoints', 'trainer_utils.load_checkpoints', (['pretrained_checkpoint_dir', 'saver2'], {}), False, 'from avod.core import trainer_utils\n')]