import tensorflow as tf
import math
import numpy as np
import random, time
import tabulate
import copy

def set_gpu_growth_or_cpu(use_cpu=False, write_info=False):

    gpus = tf.config.list_physical_devices('GPU')
    if gpus:
        if use_cpu:
            if write_info:
                print("Use CPU")
            tf.config.set_visible_devices(gpus[1:], 'GPU')
        else:
            try:
                for gpu in gpus:
                    tf.config.experimental.set_memory_growth(gpu, True)
                if write_info:
                    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
                    print(len(gpus), "Physical GPUs, ", len(logical_gpus), " Logical GPUs")
                    print('Use GPU')
            except RuntimeError as e:
                print(e)
    else:
        print("Running CPU, please check GPU drivers, CUDA, ...")

    tf.get_logger().setLevel('INFO')


def set_seed(seed, reset_session=False):
    if reset_session:
        tf.keras.backend.clear_session()
    random.seed(seed)
    np.random.seed(seed)
    tf.random.set_seed(seed)


class VerboseFitCallBack(tf.keras.callbacks.Callback):
    def __init__(self, print_lr=False):
        super(VerboseFitCallBack).__init__()
        self.columns = None
        self.st_time = 0
        self.print_lr = print_lr

    def on_epoch_begin(self, epoch, logs=None):
        self.st_time = time.time()

    def on_epoch_end(self, epoch, logs=None):
        cus_logs = copy.deepcopy(logs)
        cus_logs.pop('batch', None)
        cus_logs.pop('size', None)

        current_header = list(cus_logs.keys())
        if 'lr' in current_header:
            lr_index = current_header.index('lr')
        else:
            lr_index = len(current_header)

        if self.columns is None:

            self.columns = current_header[:lr_index] + current_header[lr_index + 1:] + ['time']
            if self.print_lr and tf.executing_eagerly():
                self.columns = ['ep', 'lr'] + self.columns
            else:
                self.columns = ['ep',] + self.columns
            # for col_index in range(len(self.columns)):
            #     if len(self.columns[col_index]) > 10:
            #         self.columns[col_index] = self.columns[col_index][:10]
        logs_values = list(cus_logs.values())

        if self.print_lr and tf.executing_eagerly():
            # Get Learning rate
            current_lr = tf.keras.backend.get_value(self.model.optimizer.lr)
            try:
                current_step = tf.cast(self.model.optimizer.iterations, tf.float32)
                current_lr = float(current_lr(current_step))
                # current_lr = tf.cast(current_lr(current_step), tf.float32)
            except:
                current_lr = float(current_lr)

        time_ep = time.time() - self.st_time
        if self.print_lr and tf.executing_eagerly():
            current_values = [epoch + 1, current_lr] + logs_values[:lr_index] + logs_values[lr_index + 1:] + [time_ep]
        else:
            current_values = [epoch + 1,] + logs_values[:lr_index] + logs_values[lr_index + 1:] + [time_ep]

        table = tabulate.tabulate([current_values], self.columns, tablefmt='simple', floatfmt='10.6g')
        if epoch % 40 == 0:
            table = table.split('\n')
            table = '\n'.join([table[1]] + table)
        else:
            table = table.split('\n')[2]
        print(table)


class CusLRScheduler(tf.keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self, initial_learning_rate, min_lr, lr_start_warmup=0., warmup_steps=10, num_constant=0, T_max=20, num_half_cycle=1.,
                 name=None):
        super(CusLRScheduler, self).__init__()
        self.initial_learning_rate = initial_learning_rate
        self.warmup_steps = warmup_steps
        self.num_constant = num_constant
        self.T_max = T_max
        self.lr_start_warmup = lr_start_warmup
        self.min_lr = min_lr
        self.num_half_cycle = num_half_cycle
        self.name = name
        pass

    def __call__(self, step):
        with tf.name_scope(self.name or "CusLRScheduler") as name:
            initial_learning_rate = tf.convert_to_tensor(self.initial_learning_rate, name="initial_learning_rate")
            dtype = initial_learning_rate.dtype

            min_lr = tf.cast(self.min_lr, dtype)
            lr_start_warmup = tf.cast(self.lr_start_warmup, dtype)

            step_cf = tf.cast(step, dtype)
            wm_steps = tf.cast(self.warmup_steps, dtype=dtype)

            warmup_ratio = tf.where(tf.less_equal(step_cf, wm_steps), step_cf / wm_steps, 0.0)
            use_warmup_lr = tf.where(tf.less_equal(step_cf, wm_steps), 1.0, 0.0)
            warmup_lr = use_warmup_lr * (lr_start_warmup + warmup_ratio * (initial_learning_rate - lr_start_warmup))

            num_constant = tf.cast(self.num_constant, dtype=dtype)

            constant_lr = tf.where(
                tf.logical_and(tf.less_equal(step_cf - wm_steps, num_constant), use_warmup_lr<1),
                initial_learning_rate, 0.0)

            t_max = tf.cast(self.T_max, dtype)
            use_consine_lr = tf.where(tf.logical_and(tf.less_equal(step_cf, t_max), tf.less(wm_steps + num_constant, step_cf)), 1.0, 0.0)
            pi_val = tf.cast(tf.constant(math.pi), dtype)
            num_half_cycle = tf.cast(self.num_half_cycle, dtype)
            cosine_lr = tf.where(use_consine_lr>0., min_lr + (initial_learning_rate - min_lr) * (1 + tf.cos(
                pi_val * num_half_cycle*(step_cf - wm_steps - num_constant) / (t_max - wm_steps - num_constant))) / 2, 0.)

            use_min_lr = tf.where(tf.less_equal(t_max, step_cf), min_lr, 0.0)

            return use_min_lr + cosine_lr + constant_lr + warmup_lr

    def get_config(self):
        ret_config = {'initial_learning_rate': self.initial_learning_rate,
                      'min_lr': self.min_lr,
                      'lr_start_warmup': self.lr_start_warmup,
                      'warmup_steps': self.warmup_steps,
                      'num_constant': self.num_constant,
                      'T_max': self.T_max,
                      'num_half_cycle': self.num_half_cycle,
                      'name': self.name}
        return ret_config

class MultiModalLoss(tf.keras.layers.Layer):
    """Adapted from https://github.com/yaringal/multi-task-learning-example"""
    def __init__(self, num_outputs=4, loss_function=tf.keras.losses.CategoricalCrossentropy, trainable=True, **kwargs):
        self.num_outputs = num_outputs
        self.loss_func = loss_function
        self.trainable = trainable
        self.config = {'num_outputs': num_outputs, 'loss_function': loss_function, 'trainable': trainable}
        super(MultiModalLoss, self).__init__(**kwargs)

    def build(self, input_shape=None):
        self.log_vars = []
        for idx in range(self.num_outputs):
            self.log_vars += [self.add_weight(name='log_var_{}'.format(idx), shape=(1, ), initializer=tf.keras.initializers.Constant(0.), trainable=self.trainable)]

        super(MultiModalLoss, self).build(input_shape)

    def multi_loss(self, ys_true, ys_pred):
        assert len(ys_true) == self.num_outputs and len(ys_pred) == self.num_outputs
        loss = 0
        for y_true, y_pred, log_var in zip(ys_true, ys_pred, self.log_vars):
            prec = tf.exp(-log_var[0])
            loss = loss + prec * self.loss_func(y_true, y_pred)

        loss = loss + tf.reduce_mean(self.log_vars)
        return loss

    def call(self, inputs):
        ys_true = inputs[: self.num_outputs]
        ys_pred = inputs[self.num_outputs: ]
        loss = self.multi_loss(ys_true, ys_pred)
        self.add_loss(loss, inputs=inputs)

        # Only return last prediction
        return inputs[-1]

    def get_config(self):
        return self.config
tensorflow.convert_to_tensor tensorflow.cast tensorflow.keras.backend.clear_session tensorflow.config.list_physical_devices tensorflow.random.set_seed tensorflow.config.experimental.set_memory_growth tensorflow.name_scope tensorflow.config.set_visible_devices tensorflow.executing_eagerly tensorflow.config.experimental.list_logical_devices tensorflow.less tensorflow.keras.backend.get_value tensorflow.less_equal tensorflow.exp tensorflow.keras.initializers.Constant tensorflow.constant tensorflow.cos numpy.random.seed tensorflow.reduce_mean tensorflow.get_logger
[(10, 'tensorflow.config.list_physical_devices', 'tf.config.list_physical_devices', (['"""GPU"""'], {}), True, 'import tensorflow as tf\n'), (35, 'random.seed', 'random.seed', (['seed'], {}), False, 'import random, time\n'), (36, 'numpy.random.seed', 'np.random.seed', (['seed'], {}), True, 'import numpy as np\n'), (37, 'tensorflow.random.set_seed', 'tf.random.set_seed', (['seed'], {}), True, 'import tensorflow as tf\n'), (34, 'tensorflow.keras.backend.clear_session', 'tf.keras.backend.clear_session', ([], {}), True, 'import tensorflow as tf\n'), (48, 'time.time', 'time.time', ([], {}), False, 'import random, time\n'), (51, 'copy.deepcopy', 'copy.deepcopy', (['logs'], {}), False, 'import copy\n'), (89, 'tabulate.tabulate', 'tabulate.tabulate', (['[current_values]', 'self.columns'], {'tablefmt': '"""simple"""', 'floatfmt': '"""10.6g"""'}), False, 'import tabulate\n'), (15, 'tensorflow.config.set_visible_devices', 'tf.config.set_visible_devices', (['gpus[1:]', '"""GPU"""'], {}), True, 'import tensorflow as tf\n'), (29, 'tensorflow.get_logger', 'tf.get_logger', ([], {}), True, 'import tensorflow as tf\n'), (73, 'tensorflow.executing_eagerly', 'tf.executing_eagerly', ([], {}), True, 'import tensorflow as tf\n'), (75, 'tensorflow.keras.backend.get_value', 'tf.keras.backend.get_value', (['self.model.optimizer.lr'], {}), True, 'import tensorflow as tf\n'), (83, 'time.time', 'time.time', ([], {}), False, 'import random, time\n'), (84, 'tensorflow.executing_eagerly', 'tf.executing_eagerly', ([], {}), True, 'import tensorflow as tf\n'), (113, 'tensorflow.name_scope', 'tf.name_scope', (["(self.name or 'CusLRScheduler')"], {}), True, 'import tensorflow as tf\n'), (114, 'tensorflow.convert_to_tensor', 'tf.convert_to_tensor', (['self.initial_learning_rate'], {'name': '"""initial_learning_rate"""'}), True, 'import tensorflow as tf\n'), (117, 'tensorflow.cast', 'tf.cast', (['self.min_lr', 'dtype'], {}), True, 'import tensorflow as tf\n'), (118, 'tensorflow.cast', 'tf.cast', (['self.lr_start_warmup', 'dtype'], {}), True, 'import tensorflow as tf\n'), (120, 'tensorflow.cast', 'tf.cast', (['step', 'dtype'], {}), True, 'import tensorflow as tf\n'), (121, 'tensorflow.cast', 'tf.cast', (['self.warmup_steps'], {'dtype': 'dtype'}), True, 'import tensorflow as tf\n'), (127, 'tensorflow.cast', 'tf.cast', (['self.num_constant'], {'dtype': 'dtype'}), True, 'import tensorflow as tf\n'), (133, 'tensorflow.cast', 'tf.cast', (['self.T_max', 'dtype'], {}), True, 'import tensorflow as tf\n'), (136, 'tensorflow.cast', 'tf.cast', (['self.num_half_cycle', 'dtype'], {}), True, 'import tensorflow as tf\n'), (175, 'tensorflow.exp', 'tf.exp', (['(-log_var[0])'], {}), True, 'import tensorflow as tf\n'), (178, 'tensorflow.reduce_mean', 'tf.reduce_mean', (['self.log_vars'], {}), True, 'import tensorflow as tf\n'), (64, 'tensorflow.executing_eagerly', 'tf.executing_eagerly', ([], {}), True, 'import tensorflow as tf\n'), (77, 'tensorflow.cast', 'tf.cast', (['self.model.optimizer.iterations', 'tf.float32'], {}), True, 'import tensorflow as tf\n'), (123, 'tensorflow.less_equal', 'tf.less_equal', (['step_cf', 'wm_steps'], {}), True, 'import tensorflow as tf\n'), (124, 'tensorflow.less_equal', 'tf.less_equal', (['step_cf', 'wm_steps'], {}), True, 'import tensorflow as tf\n'), (135, 'tensorflow.constant', 'tf.constant', (['math.pi'], {}), True, 'import tensorflow as tf\n'), (140, 'tensorflow.less_equal', 'tf.less_equal', (['t_max', 'step_cf'], {}), True, 'import tensorflow as tf\n'), (19, 'tensorflow.config.experimental.set_memory_growth', 'tf.config.experimental.set_memory_growth', (['gpu', '(True)'], {}), True, 'import tensorflow as tf\n'), (21, 'tensorflow.config.experimental.list_logical_devices', 'tf.config.experimental.list_logical_devices', (['"""GPU"""'], {}), True, 'import tensorflow as tf\n'), (130, 'tensorflow.less_equal', 'tf.less_equal', (['(step_cf - wm_steps)', 'num_constant'], {}), True, 'import tensorflow as tf\n'), (134, 'tensorflow.less_equal', 'tf.less_equal', (['step_cf', 't_max'], {}), True, 'import tensorflow as tf\n'), (134, 'tensorflow.less', 'tf.less', (['(wm_steps + num_constant)', 'step_cf'], {}), True, 'import tensorflow as tf\n'), (167, 'tensorflow.keras.initializers.Constant', 'tf.keras.initializers.Constant', (['(0.0)'], {}), True, 'import tensorflow as tf\n'), (137, 'tensorflow.cos', 'tf.cos', (['(pi_val * num_half_cycle * (step_cf - wm_steps - num_constant) / (t_max -\n    wm_steps - num_constant))'], {}), True, 'import tensorflow as tf\n')]