import abc

from mlpug.utils import get_value_at, convert_to_dict
from mlpug.trainers.callbacks.callback import Callback

import basics.base_utils as _


class LRSchedulerWrapperBase(Callback, metaclass=abc.ABCMeta):

    def __init__(self,
                 schedulers,
                 batch_level=True,
                 metric_to_monitor=None,
                 name="LRSchedulerWrapper",
                 **kwargs):
        """

        Calls given schedulers, per batch or per epoch, optionally using the current model quality
        as given by `metric_to_monitor`. The updated learning rate is added to the logs object.

        :param schedulers: A single scheduler instance or a dict or list of schedulers
        :param batch_level: True if the LR schedulers should be updated after every batch, else after every epoch
        :param metric_to_monitor: key path to metric value in the log object,
                                  e.g. `validation.window_average.perplexity`, used by the schedulers. If None, it is
                                  assumed that the schedulers don't monitor any metric.
        """
        super().__init__(name=name, **kwargs)

        schedulers = convert_to_dict("scheduler", schedulers)

        self._schedulers = schedulers
        self._batch_level = batch_level

        self._metric_to_monitor = metric_to_monitor

    @abc.abstractmethod
    def get_state(self):
        """

        :return: state, success (True or False)
        """
        self._log.error("This method is not implemented, implement it in your child class implementation")
        return None, False

    @abc.abstractmethod
    def set_state(self, state):
        """

        :param state:
        :return: success (True or False)
        """
        self._log.error("This method is not implemented, implement it in your child class implementation")
        return None, False

    def on_epoch_start(self, logs):
        if self._batch_level:
            return True

        self._init_logs(logs)

        return self._update_logs(logs)

    def on_batch_training_start(self, training_batch, logs):
        if not self._batch_level:
            return True

        self._init_logs(logs)

        return self._update_logs(logs)

    def on_batch_training_completed(self, training_batch, logs):
        if not self._batch_level:
            return True

        self._init_logs(logs)

        return self._update_lr(logs, 'global_iter')

    def on_epoch_completed(self, logs):
        if self._batch_level:
            return True

        self._init_logs(logs)

        return self._update_lr(logs, 'epoch')

    def _get_schedule_level(self):
        return 'batch' if self._batch_level else 'epoch'

    def _init_logs(self, logs):
        schedule_level = self._get_schedule_level()

        current = self._get_logs_base(logs)
        ctp = current['training_params'][schedule_level]

        if 'lr' not in ctp:
            ctp['lr'] = {}

    def _update_lr(self, logs, iter_name):
        current = self._get_logs_base(logs)

        model_quality = self._get_model_quality(current)
        training_iter = current[iter_name]

        success = False
        try:
            success = self._exec_schedulers(training_iter, model_quality)
            if not success:
                self._log.error("Updating of the learning rate(s) failed.")
        except Exception as e:
            _.log_exception(self._log, "An unexpected error occurred, "
                                       "execution of the learning rate scheduler(s) failed", e)

        return self._update_logs(logs) and success

    def _update_logs(self, logs):
        current = self._get_logs_base(logs)

        schedule_level = self._get_schedule_level()
        ctp = current['training_params'][schedule_level]

        try:
            current_lr = self._get_current_lr()

            lr = get_value_at('lr', ctp, warn_on_failure=False) or {}

            ctp['lr'] = {**lr, **current_lr}

            return True
        except Exception as e:
            _.log_exception(self._log, "An unexpected error occurred, "
                                       "unable to add current learning rate values to the logs object", e)
            return False

    def _get_model_quality(self, current_logs):
        if not self._metric_to_monitor:
            return None

        model_quality = get_value_at(self._metric_to_monitor, current_logs)

        if type(model_quality) is tuple:
            # use the first value as metric value, the other values are auxiliary results meant for other purposes
            model_quality = model_quality[0]

        return model_quality

    @abc.abstractmethod
    def _exec_schedulers(self, training_iter, model_quality=None):
        """

        :param training_iter
        :param model_quality:
        :return: On success : True
                 On failure : False
        """
        self._log.error("This method is not implemented, implement it in your child class implementation")
        return False

    @abc.abstractmethod
    def _get_current_lr(self):
        """
        :return: dict with learning rate, per optimizer
        """
        self._log.error("This method is not implemented, implement it in your child class implementation")
        return None

    def _get_callback_properties_for_hash(self):
        """
        This is used to create the unique callback hash.

        Returns a dict with properties describing the setup of the callback.
        It should at least contain the properties that influence the callback state.

        Property values should only be simple types, such as int, float, boolean and strings.
        Convert any object and function values (or similar) into a booleans
        (True = available, False = None, not available)

        :return: dict
        """
        props = super()._get_callback_properties_for_hash()
        return {
            **props,
            "batch_level": self._batch_level,
            "metric_to_monitor": self._metric_to_monitor
        }



