"""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, 25, 33, 47), 'tensorflow.Variable', 'tf.Variable', (), '', True, 'import tensorflow as tf\n'), ((81, 25, 84, 27), 'avod.builders.optimizer_builder.build', 'optimizer_builder.build', ({(82, 8, 82, 30): 'train_config.optimizer', (83, 8, 83, 24): 'global_summaries', (84, 8, 84, 26): 'global_step_tensor'}, {}), '(train_config.optimizer, global_summaries,\n    global_step_tensor)', False, 'from avod.builders import optimizer_builder\n'), ((96, 4, 96, 48), 'tensorflow.summary.scalar', 'tf.summary.scalar', ({(96, 22, 96, 37): '"""training_loss"""', (96, 39, 96, 47): 'train_op'}, {}), "('training_loss', train_op)", True, 'import tensorflow as tf\n'), ((109, 21, 115, 5), 'avod.core.summary_utils.summaries_to_keep', 'summary_utils.summaries_to_keep', (), '', False, 'from avod.core import summary_utils\n'), ((129, 19, 130, 52), 'tensorflow.summary.FileWriter', 'tf.summary.FileWriter', ({(129, 41, 129, 68): "logdir + '/' + datetime_str", (130, 41, 130, 51): 'sess.graph'}, {}), "(logdir + '/' + datetime_str, sess.graph)", True, 'import tensorflow as tf\n'), ((132, 12, 132, 77), 'tensorflow.train.Saver', 'tf.train.Saver', (), '', True, 'import tensorflow as tf\n'), ((140, 11, 140, 44), 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((177, 18, 177, 64), 'tensorflow.train.global_step', 'tf.train.global_step', ({(177, 39, 177, 43): 'sess', (177, 45, 177, 63): 'global_step_tensor'}, {}), '(sess, global_step_tensor)', True, 'import tensorflow as tf\n'), ((182, 16, 182, 27), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((45, 11, 45, 33), 'os.path.exists', 'os.path.exists', ({(45, 26, 45, 32): 'logdir'}, {}), '(logdir)', False, 'import os\n'), ((46, 8, 46, 27), 'os.makedirs', 'os.makedirs', ({(46, 20, 46, 26): 'logdir'}, {}), '(logdir)', False, 'import os\n'), ((49, 11, 49, 41), 'os.path.exists', 'os.path.exists', ({(49, 26, 49, 40): 'checkpoint_dir'}, {}), '(checkpoint_dir)', False, 'import os\n'), ((50, 8, 50, 35), 'os.makedirs', 'os.makedirs', ({(50, 20, 50, 34): 'checkpoint_dir'}, {}), '(checkpoint_dir)', False, 'import os\n'), ((70, 30, 70, 54), 'tensorflow.trainable_variables', 'tf.trainable_variables', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((87, 9, 87, 38), 'tensorflow.variable_scope', 'tf.variable_scope', ({(87, 27, 87, 37): '"""train_op"""'}, {}), "('train_op')", True, 'import tensorflow as tf\n'), ((108, 20, 108, 61), 'tensorflow.get_collection', 'tf.get_collection', ({(108, 38, 108, 60): 'tf.GraphKeys.SUMMARIES'}, {}), '(tf.GraphKeys.SUMMARIES)', True, 'import tensorflow as tf\n'), ((120, 17, 120, 33), 'tensorflow.ConfigProto', 'tf.ConfigProto', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((122, 15, 122, 40), 'tensorflow.Session', 'tf.Session', (), '', True, 'import tensorflow as tf\n'), ((124, 15, 124, 27), 'tensorflow.Session', 'tf.Session', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((127, 23, 127, 46), 'datetime.datetime.now', 'datetime.datetime.now', ({}, {}), '()', False, 'import datetime\n'), ((144, 8, 144, 60), 'avod.core.trainer_utils.load_checkpoints', 'trainer_utils.load_checkpoints', ({(144, 39, 144, 53): 'checkpoint_dir', (144, 54, 144, 59): 'saver'}, {}), '(checkpoint_dir, saver)', False, 'from avod.core import trainer_utils\n'), ((106, 26, 106, 65), 'tensorflow.contrib.memory_stats.MaxBytesInUse', 'tf.contrib.memory_stats.MaxBytesInUse', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((165, 34, 165, 58), 'tensorflow.trainable_variables', 'tf.trainable_variables', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((169, 21, 169, 65), 'tensorflow.train.Saver', 'tf.train.Saver', (), '', True, 'import tensorflow as tf\n'), ((171, 12, 171, 77), 'avod.core.trainer_utils.load_checkpoints', 'trainer_utils.load_checkpoints', ({(171, 43, 171, 68): 'pretrained_checkpoint_dir', (171, 70, 171, 76): 'saver2'}, {}), '(pretrained_checkpoint_dir, saver2)', False, 'from avod.core import trainer_utils\n'), ((187, 26, 188, 66), 'tensorflow.train.global_step', 'tf.train.global_step', ({(187, 47, 187, 51): 'sess', (188, 47, 188, 65): 'global_step_tensor'}, {}), '(sess, global_step_tensor)', True, 'import tensorflow as tf\n'), ((202, 27, 202, 38), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((152, 38, 152, 62), 'tensorflow.trainable_variables', 'tf.trainable_variables', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((156, 25, 156, 69), 'tensorflow.train.Saver', 'tf.train.Saver', (), '', True, 'import tensorflow as tf\n'), ((158, 16, 158, 81), 'avod.core.trainer_utils.load_checkpoints', 'trainer_utils.load_checkpoints', ({(158, 47, 158, 72): 'pretrained_checkpoint_dir', (158, 74, 158, 80): 'saver2'}, {}), '(pretrained_checkpoint_dir, saver2)', False, 'from avod.core import trainer_utils\n')]