import os
import datetime
import numpy as np
import tensorflow as tf

from tensorflow.keras import losses
from types import SimpleNamespace
from barbell2.utils import Logger
from autosegl3 import models
from autosegl3 import utils


def run_once(f):
    def wrapper(*args, **kwargs):
        if not wrapper.has_run:
            wrapper.has_run = True
            return f(*args, **kwargs)

    wrapper.has_run = False
    return wrapper


@run_once
def start_graph_tensorflow():
    tf.summary.trace_on(graph=True, profiler=True)


@run_once
def end_graph_tensorflow(self, log_dir):
    with self.as_default():
        tf.summary.trace_export(name="graph", step=0, profiler_outdir=log_dir)


def dice_loss(y_true, y_pred):
    smooth = 1
    union = tf.reduce_sum(y_true, axis=[1, 2, 3]) + tf.reduce_sum(y_pred, axis=[1, 2, 3])
    intersection = tf.reduce_sum(y_true * y_pred, axis=[1, 2, 3])
    score = tf.reduce_mean((2. * intersection + smooth) / (union + smooth), axis=0)
    return 1 - score


def early_stopping(loss_list, min_delta=0.005, patience=20):
    if len(list(loss_list)) // patience < 2:
        return False
    mean_previous = np.mean(loss_list[::-1][patience:2 * patience])
    mean_recent = np.mean(loss_list[::-1][:patience])
    delta_abs = np.abs(mean_recent - mean_previous)  # abs change
    delta_abs = np.abs(delta_abs / mean_previous)  # relative change
    if delta_abs < min_delta:
        print('Stopping early...')
        return True
    else:
        return False


def get_batch(images, labels, params):
    ct_batch = np.zeros(shape=[params['batch_size'],
                               params['patch_shape'][0],
                               params['patch_shape'][1],
                               params['patch_shape'][2]])
    gt_batch = np.zeros(shape=[params['batch_size'],
                               params['patch_shape'][0],
                               params['patch_shape'][1],
                               params['patch_shape'][2]])
    for patch in range(0, params['batch_size']):
        patient = random.randint(0, images.shape[-1] - 1)
        ct = images[:, :, patient]
        gt = labels[:, :, patient]
        # ct = utils.normalize(ct, 'True', params['min_bound'], params['max_bound'])
        # if random.randint(0, 1) == 1:
        #    num_augments = np.random.randint(1, params['number_of_augmentations'] + 1)
        #    ct, gt = data_augmentation.apply_augmentations(ct,
        #                                                   gt,
        #                                                   num_augments)
        ct_batch[patch, :, :, 0] = ct
        gt_batch[patch, :, :, 0] = gt
    gt_batch = tf.one_hot(np.uint8(np.squeeze(gt_batch, axis=-1)), params['num_classes'])
    return ct_batch, gt_batch


class DiceMetric(tf.keras.metrics.Metric):

    def __init__(self, name='dice_coefficient', **kwargs):
        super(DiceMetric, self).__init__(name=name, **kwargs)
        self.dice_score = self.add_weight(name='dsc', initializer='zeros')

    def update_state(self, y_true, y_pred, sample_weight=None):
        smooth = 0.000001
        union = tf.reduce_sum(y_true, axis=[1, 2, 3]) + tf.reduce_sum(y_pred, axis=[1, 2, 3])
        intersection = tf.reduce_sum(y_true * y_pred, axis=[1, 2, 3])
        score = tf.reduce_mean((2. * intersection) / (union + smooth), axis=0)
        self.dice_score.assign(score)

    def result(self):
        return self.dice_score

    def reset_states(self):
        self.dice_score.assign(0.0)


class TrainL3:

    def __init__(self, args):
        if isinstance(args, dict):
            self.params = SimpleNamespace(**args)
        else:
            self.params = args
        self.logger = Logger(prefix='trainl3', to_dir='.')
        self.train_test_split = True if self.params.test_size > 0.0 else False

    def execute(self):

        @tf.function
        def train_on_batch(im_src, gt_src):
            with tf.GradientTape() as tape:
                predictions = model(inputs=[im_src], training=True)
                regularization_loss = tf.math.add_n(model.losses)
                loss_value = loss_function(gt_src, predictions)
                total_loss = regularization_loss + loss_value
            grads = tape.gradient(total_loss, model.trainable_weights)
            optimizer_function.apply_gradients(zip(grads, model.trainable_weights))
            train_loss(total_loss)
            return predictions

        @tf.function
        def validate_on_batch(im_src, gt_src):
            predictions = model(inputs=[im_src], training=False)
            regularization_loss = tf.math.add_n(model.losses)
            loss_value = loss_function(gt_src, predictions)
            total_loss = regularization_loss + loss_value
            validation_loss(total_loss)
            return predictions

        loss_list = []
        loss_function = losses.CategoricalCrossentropy()
        optimizer_function = tf.keras.optimizers.Adam(self.params.learning_rate)

        final_model_path = None

        for iteration in range(0, self.params.num_steps + 1):

            # Define model
            model = models.unet(self.params,
                                optimizer=optimizer_function,
                                loss=loss_function)

            # Define evaluation metrics
            train_loss = tf.keras.metrics.Mean(name='train_loss')
            train_accuracy = DiceMetric()
            validation_loss = tf.keras.metrics.Mean(name='validation_loss')
            validation_accuracy = DiceMetric()

            # Create variables for various paths used for storing training information
            current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
            if not os.path.exists(self.params.output_dir):
                os.mkdir(self.params.output_dir)

            train_log_dir = self.params.output_dir + '/gradient_tape/' + current_time + '/train'
            val_log_dir = self.params.output_dir + '/gradient_tape/' + current_time + '/val'
            saved_model_path = self.params.output_dir + '/gradient_tape/' + current_time + '/saved_models/'
            saved_weights_path = self.params.output_dir + '/gradient_tape/' + current_time + '/saved_weights/'

            train_summary_writer = tf.summary.create_file_writer(train_log_dir)
            val_summary_writer = tf.summary.create_file_writer(val_log_dir)

            os.mkdir(saved_model_path)
            os.mkdir(saved_weights_path)

            images_t, labels_t, names_t = utils.load_dataset(self.params.data_path_train, self.params)
            print('Training set size: ', np.shape(images_t)[2])

            # Start training loop
            start_graph_tensorflow()

            ct_batch, gt_batch = self.get_batch(images_t, labels_t, self.params)
            train_pred = train_on_batch(ct_batch, gt_batch)

            end_graph_tensorflow(train_summary_writer, train_log_dir)

            # Evaluation step during training.
            if iteration % self.params.train_eval_step == 0:
                # Write training information to training log
                with train_summary_writer.as_default():
                    train_dice = train_accuracy(gt_batch, train_pred).numpy()
                    tf.summary.scalar('loss', train_loss.result(), step=iteration)
                    tf.summary.scalar('accuracy', train_dice, step=iteration)

                template = 'Iteration {}, Loss: {:.5}, Dice: {:.5}'
                print(template.format(iteration + 1,
                                      train_loss.result(),
                                      train_dice))

            # Evaluation step for validation.
            # TODO: Figure out best way to call train.py with train/test split or not
            if self.train_test_split:
                images_v, labels_v, names_v = utils.load_dataset(self.params.data_path_val, self.params)
                print('Test set size: ', np.shape(images_v)[2])
                if iteration % self.params.val_eval_step == 0:
                    ct_batch_val, gt_batch_val = self.get_batch(images_v, labels_v, self.params)

                    # gt_batch_val = tf.one_hot(np.uint8(np.squeeze(gt_batch_val)), params.dict['num_classes'])
                    val_pred = validate_on_batch(ct_batch_val, gt_batch_val)

                    # Write validation information to log
                    with val_summary_writer.as_default():
                        validation_dice = validation_accuracy(gt_batch_val, val_pred).numpy()
                        tf.summary.scalar('loss', validation_loss.result(), step=iteration)
                        tf.summary.scalar('accuracy', validation_dice, step=iteration)
                        loss_list.append(validation_loss.result())

                    template = 'Iteration {}, Validation Loss: {:.5}, Validation Dice: {:.5}'
                    print(template.format(iteration + 1,
                                          validation_loss.result(),
                                          validation_dice))

                    # Earling stopping when loss in the past 'patience' train_eval_steps
                    # is smaller than 'min_delta'. Breaks loop.
                    early_stop = self.early_stopping(loss_list, min_delta=0.001, patience=10)
                    if early_stop:
                        print("Early stopping signal received at iteration = %d/%d" % (iteration, self.params.num_steps))
                        print("Terminating training ")
                        model.save(os.path.join(saved_model_path, 'model_' + str(iteration)))
                        model.save_weights(os.path.join(saved_weights_path, 'model_weights' + str(iteration) + '.h5'))
                        final_model = model
                        final_model_path = os.path.join(saved_model_path, 'final_model')
                        final_model.save(final_model_path)
                        final_model.save_weights(os.path.join(saved_weights_path, 'final_model_weights' + '.h5'))
                        break

            # Save the model at predefined step numbers.
            if iteration % self.params.save_model_step == 0:
                model.save(os.path.join(saved_model_path, 'model_' + str(iteration)))
                model.save_weights(os.path.join(saved_weights_path, 'model_weights' + str(iteration) + '.h5'))

        return final_model_path
