import tensorflow as tf
import matplotlib.pyplot as plt
import os
import utils
import re
import numpy as np
from tensorflow.keras import backend

#######################################################################################################################
class monitor(tf.keras.callbacks.Callback):
    def __init__(self, save_dir, dataset=None):
        super(monitor, self).__init__()
        self.save_dir = save_dir
        self.dataset = dataset

    #
    def on_epoch_end(self, epoch, logs=None):
        epoch += self.params['initial_epoch'] + 1
        ncols, nrows = 3, 5
        size_rate = 3
        figure, axs = plt.subplots(ncols, nrows, figsize=(nrows * size_rate, ncols * size_rate))
        figure.suptitle(f'Epoch: "{epoch}"', fontsize=12)
        figure.tight_layout()
        for cn, an in self.dataset:
            mask = self.model.predict_on_batch(an)
            an, mask = tf.squeeze(an), tf.squeeze(mask)
            for r in range(nrows):
                slice = int(an.shape[2] * 0.4) + (r * 2)
                axs[0][r].imshow(an[:, slice, :], cmap='gray')
                axs[1][r].imshow(mask[:, slice, :], cmap='gray')
                # axs[0][r].imshow((mask-an)[slice, :, :], cmap='Reds', alpha=0.4)
                axs[2][r].imshow((mask-an)[:, slice, :], cmap='gray')
                # axs[1][r].imshow((mask-an)[:, slice, :], cmap='Reds', alpha=0.4)
                # axs[2][r].imshow(mask[:, :, slice], cmap='gray')
                # axs[2][r].imshow((mask-an)[:, :, slice], cmap='Reds', alpha=0.4)
        save_path = os.path.join(self.save_dir, f'{epoch:05}.png')
        plt.savefig(save_path, dpi=150)
        plt.close('all')

#######################################################################################################################
class lr_scheduler(tf.keras.callbacks.Callback):
    def __init__(self):
        super(lr_scheduler, self).__init__()

    def on_epoch_end(self, epoch, logs=None):
        if not hasattr(self.model.G_optimizer, 'lr'):
            raise ValueError('Optimizer must have a "lr" attribute.')
        lr = float(backend.get_value(self.model.G_optimizer.lr))
        lr *= tf.math.exp(-0.1)
        if not isinstance(lr, (tf.Tensor, float, np.float32, np.float64)):
            raise ValueError('The output of the "schedule" function '
                             f'should be float. Got: {lr}')
        if isinstance(lr, tf.Tensor) and not lr.dtype.is_floating:
            raise ValueError(
                f'The dtype of `lr` Tensor should be float. Got: {lr.dtype}')
        backend.set_value(self.model.G_optimizer.lr, backend.get_value(lr))
        backend.set_value(self.model.D_optimizer.lr, backend.get_value(lr))

        epoch += self.params['initial_epoch'] + 1
        print(f'\n[Callback] Epoch {epoch + 1}: setting learning to "{lr}".')
        print(f'd_lr: "{float(backend.get_value(self.model.D_optimizer.lr))}"')
        print(f'g_lr: "{float(backend.get_value(self.model.G_optimizer.lr))}"\n')

#######################################################################################################################
class ckpt(tf.keras.callbacks.Callback):
    def __init__(self, save_dir):
        super(ckpt, self).__init__()
        self.save_dir = save_dir
        # self.manager = tf.train.CheckpointManager(self.ckpt, directory=save_dir, max_to_keep=10)

    def on_train_begin(self, logs=None):
        self.ckpt = tf.train.Checkpoint(
            D_optimizer=self.model.D_optimizer,
            G_optimizer=self.model.G_optimizer,
            D=self.model.disc,
            G=self.model.gen,
        )
        latest_ckpt = tf.train.latest_checkpoint(self.save_dir)
        if latest_ckpt:
            self.initial_epoch = int(re.findall(r"EP_(\d+)-", latest_ckpt)[0])
            state = self.ckpt.restore(latest_ckpt)
            """
            - Raises an exception if any variables are unmatched
            - Raises an exception if any existing Python objects in the dependency graph are unmatched.
            - Asserts that something aside from the root object was matched
            """
            # state.assert_consumed()
            # state.assert_existing_objects_matched()
            state.assert_nontrivial_match()
            print(f'\n[Model] Restored from Epoch "{self.initial_epoch}".\n')
        else:
            print(f'\n[Model] There is no "latest_checkpoint", Train from scratch !\n')
            self.initial_epoch = 0

        ### finally
        self.params['initial_epoch'] = self.initial_epoch


    #
    def on_epoch_end(self, epoch, logs=None):
        epoch += self.initial_epoch + 1
        file_prefix = f'EP_{epoch}'
        save_file_path = utils.join_dir([self.save_dir, file_prefix])
        self.ckpt.save(save_file_path)
        print(f'\n[Model] Save model, epoch:"{epoch}"\n')

#######################################################################################################################
class set_lr(tf.keras.callbacks.Callback):
    def __init__(self, lr):
        super(set_lr, self).__init__()
        self.lr = lr

    def on_train_begin(self, logs=None):
        if not isinstance(self.lr, (tf.Tensor, float, np.float32, np.float64)):
            raise ValueError(f'LR should be float. Got: {self.lr}')
        if isinstance(self.lr, tf.Tensor) and not self.lr.dtype.is_floating:
            raise ValueError(
                f'The dtype of `lr` Tensor should be float. Got: {self.lr.dtype}')
        backend.set_value(self.model.G_optimizer.lr, backend.get_value(self.lr))
        backend.set_value(self.model.D_optimizer.lr, backend.get_value(self.lr))
#######################################################################################################################