import time
import types
import logging
from typing import List, Dict, Optional, Callable
from collections import defaultdict

import tensorflow as tf
import tensorflow.keras as ks


# == CUSTOM ACTIVATION FUNCTIONS ============================================================================
# ===========================================================================================================

def shifted_sigmoid(x, multiplier: float = 1, shift: float = 10) -> float:
    return ks.backend.sigmoid(multiplier * (x - shift))


# == TRAINING SEGMENTATION ==================================================================================
# ===========================================================================================================


class RecompilableMixin:

    def __init__(self):
        self.compile_args: Optional[list] = None
        self.compile_kwargs: Optional[dict] = None

    def compile(self, *args, **kwargs):
        ks.models.Model.compile(self, *args, **kwargs)
        self.compile_args = args
        self.compile_kwargs = kwargs

    def recompile(self):
        self.compile(*self.compile_args, **self.compile_kwargs)


class SegmentedFitProcess:

    def __init__(self,
                 model: ks.models.Model,
                 fit_kwargs: dict):
        self.model = model
        self.fit_kwargs = fit_kwargs

        self.total_epochs = fit_kwargs['epochs']
        self.callbacks = dict()
        self.callbacks[self.total_epochs] = lambda model, history: None
        self.current_epoch = 0

    def __call__(self):
        """
        This method actually executes the various model.fit() operations which are part of the overall
        segmented process. It then returns a ks History object which contains the *overall* history which
        consists of the merged histories of all the individual processes.

        :return:
        """
        hists = []

        # self.callbacks: Dict[int, Callable] is a dictionary whose keys are integer epoch numbers which
        # are generally between 0 and the total number of epochs defined for the entire fit process.
        # The corresponding values are callbacks which are applied *on the model* at the point at which the
        # overall process hits that epoch.
        for epoch, cb in sorted(self.callbacks.items(), key=lambda i: i[0]):
            # The next fit process has to have as many epochs as starting from the current one and until the
            # next callback has to be applied.
            epoch_diff = epoch - self.current_epoch

            # For each individual fit process we can just reuse the majority of the fit kwargs, but we need
            # to replace the epoch count with the local count only until the next callback is due.
            fit_kwargs = self.fit_kwargs.copy()
            fit_kwargs['epochs'] = epoch_diff

            # We train the model and then after it is done we apply the callback.
            hist = self.model.fit(**fit_kwargs)
            hists.append(hist)
            cb(
                model=self.model,
                history=hist
            )

            self.current_epoch += epoch_diff

        merged_hist = self.merge_histories(hists)
        return merged_hist

    def merge_histories(self, hists: List[ks.callbacks.History]):
        merged_hist = ks.callbacks.History()

        merged_history = defaultdict(list)
        for hist in hists:
            for key, value in hist.history.items():
                if isinstance(value, list):
                    merged_history[key] += value

        merged_hist.history = merged_history
        return merged_hist

    def add_callback(self,
                     epochs: int,
                     callback: Callable[[ks.models.Model, ks.callbacks.History], None]
                     ) -> None:
        self.callbacks[epochs] = callback

# This section is a lot of black magic
# ------------------------------------

def epochs_once(epoch: int):

    def register(self, instance, fit_process, fit_kwargs):
        # https://stackoverflow.com/questions/1015307/python-bind-an-unbound-method
        bound_method = self.__get__(instance, instance.__class__)
        fit_process.add_callback(epoch, bound_method)

    def local(func: Callable):
        method = types.MethodType(register, func)
        setattr(func, 'register', method)
        setattr(func, 'is_epoch_callback', True)

        return func

    return local


class FitManager:

    def __init__(self):
        pass

    def __call__(self,
                 model: ks.models.Model,
                 fit_kwargs: dict) -> SegmentedFitProcess:
        fit_process = SegmentedFitProcess(model, fit_kwargs)
        # self.register_callbacks(fit_process, fit_kwargs, model)
        callbacks = self.gather_callbacks()
        for callback in callbacks:
            callback.register(self, fit_process, fit_kwargs)

        return fit_process

    def gather_callbacks(self) -> List[Callable]:
        callbacks = []
        for name in dir(self):
            attribute = getattr(self, name)
            if callable(attribute) and hasattr(attribute, 'is_epoch_callback'):
                callbacks.append(attribute)

        return callbacks

    def register_callbacks(self,
                           fit_process: SegmentedFitProcess,
                           fit_kwargs: dict,
                           model: ks.models.Model) -> None:
        return


class LockExplanationManager(FitManager):

    def __init__(self,
                 epoch_threshold: int):
        FitManager.__init__(self)
        self.epoch_threshold = epoch_threshold

    def register_callbacks(self, fit_process, fit_kwargs, model):
        fit_process.add_callback(self.epoch_threshold, self.lock_model_explanation_parameters)

    def lock_model_explanation_parameters(self, model, history):
        model.lock_explanation_parameters()
        model.recompile()


class NoLoss(ks.losses.Loss):

    def __init__(self,
                 *args,
                 name='no_loss',
                 **kwargs):
        super(NoLoss, self).__init__(*args, name=name, **kwargs)

    def call(self, y_true, y_pred):
        return 0


def bce(y_true, y_pred):
    y_true = tf.cast(y_true, tf.float32)
    loss = y_true * tf.math.log(y_pred + 1e-7)
    loss += (1 - y_true) * tf.math.log(1 - y_pred + 1e-7)
    loss = -loss
    return tf.reduce_mean(loss)


def mae(y_true, y_pred):
    y_true = tf.cast(y_true, tf.float32)
    nan_multiplier = tf.where(tf.math.is_nan(y_true), 0., 1.)

    y_true = tf.where(tf.math.is_nan(y_true), 0., y_true)
    loss = tf.abs(y_true - y_pred)

    loss *= nan_multiplier
    loss = tf.reduce_sum(loss, axis=-1)

    return tf.reduce_mean(loss)


mae.name = 'mean_absolute_error'


def mse(y_true, y_pred):
    nan_multiplier = tf.where(tf.math.is_nan(y_true), 0., 1.)

    y_true = tf.cast(y_true, tf.float32)
    y_true = tf.where(tf.math.is_nan(y_true), 0., y_true)
    loss = tf.square(y_true - y_pred)

    loss *= nan_multiplier
    loss = tf.reduce_sum(loss, axis=-1)
    # print(loss)

    return tf.reduce_mean(loss)


mse.name = 'mean_squared_error'


class ExplanationLoss(ks.losses.Loss):

    def __init__(self,
                 loss_function: ks.losses.Loss = bce,
                 mask_empty_explanations: bool = False,
                 reduce: bool = False,
                 factor: float = 1):
        ks.losses.Loss.__init__(self)

        self.loss_function = loss_function
        self.mask_empty_explanations = mask_empty_explanations
        self.factor = factor
        self.reduce = reduce
        self.name = 'explanation_loss'

    def call(self, y_true, y_pred):
        loss = self.loss_function(y_true, y_pred)

        if self.mask_empty_explanations:
            mask = tf.cast(tf.reduce_max(y_true, axis=-1) > 0, dtype=tf.float32)
            loss *= mask

        if self.reduce:
            loss = tf.reduce_mean(loss, axis=-1)

        return self.factor * loss


class EpochCounterCallback(ks.callbacks.Callback):

    def __init__(self):
        super(EpochCounterCallback, self).__init__()
        self.epoch = 0

    def on_epoch_begin(self, epoch, logs=None):
        self.epoch += 1


class LogProgressCallback(EpochCounterCallback):

    def __init__(self,
                 logger: logging.Logger,
                 identifier: str,
                 epoch_step: int):
        super(LogProgressCallback, self).__init__()

        self.logger = logger
        self.identifier = identifier
        self.epoch_step = epoch_step

        self.start_time = time.time()
        self.elapsed_time = 0

    def on_epoch_end(self, epoch, logs=None):
        if (self.epoch % self.epoch_step == 0 or self.epoch == 1) and logs is not None:
            self.elapsed_time = time.time() - self.start_time
            value = logs[self.identifier]

            message_parts = [
                f'   epoch {str(self.epoch):<5}: ',
                f'{self.identifier}={value:.3f} ',
                f'elapsed_time={self.elapsed_time:.1f}s'
            ]
            if 'output_1_loss' in logs:
                message_parts.append(f'output_1_loss={logs["output_1_loss"]:.3f} ')
            if 'output_2_loss' in logs:
                message_parts.append(f'output_2_loss={logs["output_2_loss"]:.3f}')
            if 'output_3_loss' in logs:
                message_parts.append(f'output_3_loss={logs["output_3_loss"]:.3f}')
            if 'loss' in logs:
                message_parts.append(f'total_training_loss={logs["loss"]:.3f} ')
            if 'exp_loss' in logs:
                message_parts.append(f'exp_loss={logs["exp_loss"]:.3f}')

            self.logger.info(' '.join(message_parts))

    def reset(self):
        self.epoch = 0
        self.start_time = time.time()
        self.elapsed_time = 0