# -*- coding:utf-8 -*-
# Author: Xue Yang <yangxue-2019-sjtu@sjtu.edu.cn>, <yangxue0827@126.com>
# License: Apache-2.0 license
# Copyright (c) SJTU. ALL rights reserved.

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

import os
import sys
import time

import numpy as np
import tensorflow as tf

sys.path.append("../")

from dataloader.dataset.read_tfrecord import ReadTFRecord
from alpharotate.libs.utils.show_box_in_tensor import DrawBoxTensor
from alpharotate.utils import tools


class Train(object):
    def __init__(self, cfgs):
        self.cfgs = cfgs
        self.reader = ReadTFRecord(cfgs)
        self.drawer = DrawBoxTensor(cfgs)

    def stats_graph(self, graph):
        """
        :param graph:
        :return: Flops (may be wrong) and params
        """
        flops = tf.profiler.profile(graph, options=tf.profiler.ProfileOptionBuilder.float_operation())
        params = tf.profiler.profile(graph, options=tf.profiler.ProfileOptionBuilder.trainable_variables_parameter())
        print('FLOPs: {};    Trainable params: {}'.format(flops.total_float_ops, params.total_parameters))

    def average_gradients(self, tower_grads):
        """Calculate the average gradient for each shared variable across all towers.
        Note that this function provides a synchronization point across all towers.
        Args:
          tower_grads: List of lists of (gradient, variable) tuples. The outer list
            is over individual gradients. The inner list is over the gradient
            calculation for each tower.
        Returns:
           List of pairs of (gradient, variable) where the gradient has been averaged
           across all towers.
        """
        average_grads = []
        for grad_and_vars in zip(*tower_grads):
            # Note that each grad_and_vars looks like the following:
            #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
            grads = []
            for g, _ in grad_and_vars:
                # Add 0 dimension to the gradients to represent the tower.
                expanded_g = tf.expand_dims(g, 0)

                # Append on a 'tower' dimension which we will average over below.
                grads.append(expanded_g)

            # Average over the 'tower' dimension.
            grad = tf.concat(axis=0, values=grads)
            grad = tf.reduce_mean(grad, 0)

            # Keep in mind that the Variables are redundant because they are shared
            # across towers. So .. we will just return the first tower's pointer to
            # the Variable.
            v = grad_and_vars[0][1]
            grad_and_var = (grad, v)
            average_grads.append(grad_and_var)
        return average_grads

    def sum_gradients(self, tower_grads):
        """Calculate the average gradient for each shared variable across all towers.
        Note that this function provides a synchronization point across all towers.
        Args:
          tower_grads: List of lists of (gradient, variable) tuples. The outer list
            is over individual gradients. The inner list is over the gradient
            calculation for each tower.
        Returns:
           List of pairs of (gradient, variable) where the gradient has been averaged
           across all towers.
        """
        sum_grads = []
        for grad_and_vars in zip(*tower_grads):
            # Note that each grad_and_vars looks like the following:
            #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
            grads = []
            for g, _ in grad_and_vars:
                # Add 0 dimension to the gradients to represent the tower.
                expanded_g = tf.expand_dims(g, 0)

                # Append on a 'tower' dimension which we will average over below.
                grads.append(expanded_g)

            # Average over the 'tower' dimension.
            grad = tf.concat(axis=0, values=grads)
            grad = tf.reduce_sum(grad, 0)

            # Keep in mind that the Variables are redundant because they are shared
            # across towers. So .. we will just return the first tower's pointer to
            # the Variable.
            v = grad_and_vars[0][1]
            grad_and_var = (grad, v)
            sum_grads.append(grad_and_var)
        return sum_grads

    def warmup_lr(self, init_lr, global_step, warmup_step, num_gpu):
        def warmup(end_lr, global_step, warmup_step):
            start_lr = end_lr * 0.1
            global_step = tf.cast(global_step, tf.float32)
            return start_lr + (end_lr - start_lr) * global_step / warmup_step

        def decay(start_lr, global_step, num_gpu):
            lr = tf.train.piecewise_constant(global_step,
                                             boundaries=[np.int64(self.cfgs.DECAY_STEP[0] // num_gpu),
                                                         np.int64(self.cfgs.DECAY_STEP[1] // num_gpu),
                                                         np.int64(self.cfgs.DECAY_STEP[2] // num_gpu)],
                                             values=[start_lr, start_lr / 10., start_lr / 100., start_lr / 1000.])
            return lr

        return tf.cond(tf.less_equal(global_step, warmup_step),
                       true_fn=lambda: warmup(init_lr, global_step, warmup_step),
                       false_fn=lambda: decay(init_lr, global_step, num_gpu))

    def warmup_and_cosine_lr(self, init_lr, global_step, warmup_step, decay_steps, alpha=1e-6):
        def warmup(end_lr, global_step, warmup_step):
            start_lr = end_lr * 0.1
            global_step = tf.cast(global_step, tf.float32)
            return start_lr + (end_lr - start_lr) * global_step / warmup_step

        def cosine_lr(init_lr, global_step, decay_steps, alpha=0.0):
            return tf.train.cosine_decay(learning_rate=init_lr,
                                         global_step=global_step - warmup_step,
                                         decay_steps=decay_steps - warmup_step,
                                         alpha=alpha)

        return tf.cond(tf.less_equal(global_step, warmup_step),
                       true_fn=lambda: warmup(init_lr, global_step, warmup_step),
                       false_fn=lambda: cosine_lr(init_lr, global_step, decay_steps, alpha))

    def loss_dict(self, inputs, num_gpu):
        total_loss_dict = {'total_losses': tf.constant(0., tf.float32)}
        total_losses = 0.0
        for k in inputs.keys():
            if k not in total_loss_dict.keys():
                total_loss_dict[k] = tf.constant(0., tf.float32)
            total_losses += inputs[k]
            total_loss_dict[k] += inputs[k] / num_gpu

        total_losses /= num_gpu
        total_loss_dict['total_losses'] += total_losses
        return total_loss_dict, total_losses

    def log_printer(self, deter, optimizer, global_step, tower_grads, total_loss_dict, num_gpu, graph):
        for k in total_loss_dict.keys():
            tf.summary.scalar('{}/{}'.format(k.split('_')[0], k), total_loss_dict[k])

        if len(tower_grads) > 1:
            grads = self.sum_gradients(tower_grads)
        else:
            grads = tower_grads[0]

        if self.cfgs.MUTILPY_BIAS_GRADIENT is not None:
            final_gvs = []
            with tf.variable_scope('Gradient_Mult'):
                for grad, var in grads:
                    scale = 1.
                    if '/biases:' in var.name:
                        scale *= self.cfgs.MUTILPY_BIAS_GRADIENT
                    if 'conv_new' in var.name:
                        scale *= 3.
                    if not np.allclose(scale, 1.0):
                        grad = tf.multiply(grad, scale)

                    final_gvs.append((grad, var))
            apply_gradient_op = optimizer.apply_gradients(final_gvs, global_step=global_step)
        else:
            apply_gradient_op = optimizer.apply_gradients(grads, global_step=global_step)

        variable_averages = tf.train.ExponentialMovingAverage(0.9999, global_step)
        variables_averages_op = variable_averages.apply(tf.trainable_variables())

        train_op = tf.group(apply_gradient_op, variables_averages_op)
        # train_op = optimizer.apply_gradients(final_gvs, global_step=global_step)
        summary_op = tf.summary.merge_all()

        restorer, restore_ckpt = deter.get_restorer()
        saver = tf.train.Saver(max_to_keep=2)

        init_op = tf.group(
            tf.global_variables_initializer(),
            tf.local_variables_initializer()
        )

        tfconfig = tf.ConfigProto(
            allow_soft_placement=True, log_device_placement=False)
        tfconfig.gpu_options.allow_growth = True
        with tf.Session(config=tfconfig) as sess:
            sess.run(init_op)

            # sess.run(tf.initialize_all_variables())
            coord = tf.train.Coordinator()
            threads = tf.train.start_queue_runners(coord=coord, sess=sess)

            summary_path = os.path.join(self.cfgs.SUMMARY_PATH, self.cfgs.VERSION)
            tools.makedirs(summary_path)
            summary_writer = tf.summary.FileWriter(summary_path, graph=sess.graph)

            if not restorer is None:
                restorer.restore(sess, restore_ckpt)
                print('restore model')

            self.stats_graph(graph)

            for step in range(self.cfgs.MAX_ITERATION // num_gpu):
                training_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))

                if step % self.cfgs.SHOW_TRAIN_INFO_INTE != 0 and step % self.cfgs.SMRY_ITER != 0:
                    _, global_stepnp = sess.run([train_op, global_step])

                else:
                    if step % self.cfgs.SHOW_TRAIN_INFO_INTE == 0 and step % self.cfgs.SMRY_ITER != 0:
                        start = time.time()

                        _, global_stepnp, total_loss_dict_ = \
                            sess.run([train_op, global_step, total_loss_dict])

                        end = time.time()

                        print('***'*24)
                        print("%s: global_step:%d  current_step:%d"
                              % (training_time, (global_stepnp-1)*num_gpu, step*num_gpu))

                        seconds = (self.cfgs.MAX_ITERATION - (global_stepnp-1) * num_gpu) * (end - start) / num_gpu
                        m, s = divmod(seconds, 60)
                        h, m = divmod(m, 60)
                        d, h = divmod(h, 24)
                        print("speed: %.3fs, remaining training time: %02d:%02d:%02d:%02d" % ((end - start) / num_gpu, d, h, m, s))

                        loss_str = ''
                        for k in total_loss_dict_.keys():
                            loss_str += '%s:%.3f\n' % (k, total_loss_dict_[k])
                        print(loss_str)

                        if np.isnan(total_loss_dict_['total_losses']):
                            sys.exit(0)

                    else:
                        if step % self.cfgs.SMRY_ITER == 0:
                            _, global_stepnp, summary_str = sess.run([train_op, global_step, summary_op])
                            summary_writer.add_summary(summary_str, (global_stepnp-1)*num_gpu)
                            summary_writer.flush()

                if (step > 0 and step % (self.cfgs.SAVE_WEIGHTS_INTE // num_gpu) == 0) or (step >= (self.cfgs.MAX_ITERATION // num_gpu - 1)) or (global_stepnp-1)*num_gpu > self.cfgs.MAX_ITERATION:

                    save_dir = os.path.join(self.cfgs.TRAINED_CKPT, self.cfgs.VERSION)
                    if not os.path.exists(save_dir):
                        os.makedirs(save_dir)

                    save_ckpt = os.path.join(save_dir, '{}_'.format(self.cfgs.DATASET_NAME) +
                                             str((global_stepnp-1)*num_gpu) + 'model.ckpt')
                    saver.save(sess, save_ckpt)
                    print('Weights had been saved')

                if (global_stepnp-1)*num_gpu > self.cfgs.MAX_ITERATION:
                    break

            print('***' * 24)
            print('End of training.')

            coord.request_stop()
            coord.join(threads)












