"""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', True, 'import tensorflow as tf\n'), (81, 'avod.builders.optimizer_builder.build', 'optimizer_builder.build', False, 'from avod.builders import optimizer_builder\n'), (96, 'tensorflow.summary.scalar', 'tf.summary.scalar', True, 'import tensorflow as tf\n'), (109, 'avod.core.summary_utils.summaries_to_keep', 'summary_utils.summaries_to_keep', False, 'from avod.core import summary_utils\n'), (129, 'tensorflow.summary.FileWriter', 'tf.summary.FileWriter', True, 'import tensorflow as tf\n'), (132, 'tensorflow.train.Saver', 'tf.train.Saver', 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', True, 'import tensorflow as tf\n'), (182, 'time.time', 'time.time', False, 'import time\n'), (45, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (46, 'os.makedirs', 'os.makedirs', False, 'import os\n'), (49, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (50, 'os.makedirs', 'os.makedirs', False, 'import os\n'), (70, 'tensorflow.trainable_variables', 'tf.trainable_variables', True, 'import tensorflow as tf\n'), (87, 'tensorflow.variable_scope', 'tf.variable_scope', True, 'import tensorflow as tf\n'), (108, 'tensorflow.get_collection', 'tf.get_collection', True, 'import tensorflow as tf\n'), (120, 'tensorflow.ConfigProto', 'tf.ConfigProto', True, 'import tensorflow as tf\n'), (122, 'tensorflow.Session', 'tf.Session', 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', 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', True, 'import tensorflow as tf\n'), (171, 'avod.core.trainer_utils.load_checkpoints', 'trainer_utils.load_checkpoints', False, 'from avod.core import trainer_utils\n'), (187, 'tensorflow.train.global_step', 'tf.train.global_step', 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', True, 'import tensorflow as tf\n'), (158, 'avod.core.trainer_utils.load_checkpoints', 'trainer_utils.load_checkpoints', False, 'from avod.core import trainer_utils\n')]