from enum import Enum

import math

from mlpug.trainers.callbacks.callback import Callback
from mlpug.mlpug_exceptions import CallbackInvalidException
from mlpug.utils import has_key, get_key_paths, SlidingWindow

from mlpug.evaluation import MetricEvaluatorBase

import basics.base_utils as _

from mlpug.utils import get_value_at, set_value_at


class MetricsLoggingMode(Enum):
    # TODO : Add BATCH_METRICS_ONLY Option
    BATCH_AND_WINDOW_AVERAGE_METRICS = "BATCH_AND_WINDOW_AVERAGE_METRICS"
    WINDOW_AVERAGE_METRICS = "WINDOW_AVERAGE_METRICS"
    WHOLE_DATASET_METRICS = "WHOLE_DATASET_METRICS"

    @staticmethod
    def will_log_window_average_metrics(mode):
        return mode in {MetricsLoggingMode.BATCH_AND_WINDOW_AVERAGE_METRICS,
                        MetricsLoggingMode.WINDOW_AVERAGE_METRICS}

    def __str__(self):
        return str(self.value)


class MetricsLoggerBase(Callback):

    def __init__(self,
                 metric_evaluator,
                 dataset_name,
                 batch_level,
                 logging_mode,
                 dataset=None,
                 evaluate_settings=None,
                 batch_averaging_window=None,
                 log_condition_func=None,
                 sliding_window_factory=None,
                 name=None,
                 **kwargs):
        """

        Different modes of operation:

        Batch level: After every batch training iteration:
            Calculate metrics over a sliding window
                log metrics per batch and over window   batch_level = True, mode = BATCH_AND_WINDOW_AVERAGE_METRICS
                log metrics over window only            batch_level = True, mode = WINDOW_AVERAGE_METRICS
            Calculate metrics over whole data set       batch_level = True, mode = WHOLE_DATASET_METRICS

        Epoch level: After every training epoch:
            Calculate metrics over whole data set       batch_level = False, mode = WHOLE_DATASET_METRICS


        `log_condition_func` can be used to control if metrics are calculated

        :param dataset_name:
        :type dataset_name:

        :param batch_level: See general description above
        :type batch_level:

        :param logging_mode: See general description above
        :type logging_mode:

        :param dataset: Dataset, where the items are batches, to calculate the metrics with
        :type dataset:

        :param metric_evaluator: MetricEvaluator instance, or other object with the same public interface
        :type metric_evaluator:

        :param evaluate_settings: Default evaluation settings to use, when no evaluate_settings are available in the
                                  logs object provided to the callback methods
        :type evaluate_settings:

        :param batch_averaging_window: Window length in batches, over which the average metric must be calculated
        :type batch_averaging_window: Int

        :param sliding_window_factory: Optional. Function to create SlidingWindow instances

        :param name:
        :type name:
        """
        if name is None:
            name = self.__class__.__name__
            name = f"{name}[{dataset_name}]"

        super().__init__(name=name, **kwargs)

        self._dataset = dataset
        self._dataset_name = dataset_name

        self._metric_evaluator = metric_evaluator

        self._batch_level = batch_level

        self._logging_mode = logging_mode

        self._evaluate_settings = evaluate_settings

        self._batch_averaging_window = batch_averaging_window

        self._log_condition_func = log_condition_func or (lambda logs, dataset_batch: True)

        self._sliding_window_factory = SlidingWindow if not callable(sliding_window_factory) else sliding_window_factory

        self._name = name

        self._metric_windows = {}

        self._validate()

        if not self.instance_valid():
            raise CallbackInvalidException(self._name)

    def get_state(self):
        """

        :return: state, success (True or False)
        """
        return {
                    "metric_windows": {metric_path: window.get_state()
                                       for metric_path, window in self._metric_windows.items()}
               }, True

    def set_state(self, state):
        """

        :param state:
        :return: success (True or False)
        """

        if not self._check_state(state):
            self._log.error('Given state is invalid, unable to set state')
            return False

        self._metric_windows = {metric_path: self._sliding_window_factory(state=window_state)
                                for metric_path, window_state in state["metric_windows"].items()}

        return True

    def on_training_start(self,
                          num_epochs,
                          num_batches_per_epoch,
                          start_epoch,
                          start_batch,
                          start_update_iter):

        will_log_window_average_metrics = MetricsLoggingMode.will_log_window_average_metrics(self._logging_mode)

        if not self.instance_valid():
            if will_log_window_average_metrics:
                self._log.error(f"{self} is not valid, unable to set up averaging window ... ")
            else:
                self._log.error(f"{self} is not valid, skipping this hook ... ")
            return False

        if will_log_window_average_metrics and \
                (self._batch_averaging_window is None or self._batch_averaging_window <= 0):
            self._log.error(f"A valid batch processing window is required for "
                            f"metric logging mode ({self._logging_mode}, the {self} will not function")
            self._valid = False

            return False

        if will_log_window_average_metrics:
            self._init_metric_windows()

        return True

    def on_batch_training_completed(self, dataset_batch, logs):
        if not self.instance_valid():
            self._log.error(f"{self} is not valid, skipping this hook ... ")
            return False

        if not self._batch_level:
            return True

        if not self._log_condition_func(logs=logs, dataset_batch=dataset_batch):
            return True

        self._init_logs(logs)

        if self._logging_mode is MetricsLoggingMode.WHOLE_DATASET_METRICS:
            return self._calc_whole_dataset_metrics(logs, f"{self._dataset_name}.dataset")
        else:
            current = self._get_logs_base(logs)

            batch_metrics = {}
            if not self._calc_batch_metric_data_from(dataset_batch, batch_metrics, logs):
                return False

            base_path = f"{self._dataset_name}.batch"
            dataset_batch_logs = get_value_at(base_path, current)

            # Merge in new batch level results
            dataset_batch_logs = {**dataset_batch_logs, **batch_metrics}
            if self._logging_mode is MetricsLoggingMode.BATCH_AND_WINDOW_AVERAGE_METRICS:
                set_value_at(base_path, current, dataset_batch_logs)

            metric_names = self._metric_evaluator.get_metric_names()
            metric_paths = get_key_paths(dataset_batch_logs,
                                         keys_to_consider=metric_names,
                                         keys_not_to_consider=["auxiliary_results"])

            self._update_metrics_windows_for(metric_paths, dataset_batch_logs, base_path=base_path)

            # gather all window data
            batch_metrics_lists = {p: s.window for p, s in self._metric_windows.items()}
            if not self._reduce(batch_metrics_lists, current[self._dataset_name]['window_average']):
                return False

            return True

    def on_epoch_completed(self, logs):
        if not self.instance_valid():
            self._log.error(f"{self} is not valid, skipping this hook ... ")
            return False

        if self._batch_level:
            return True

        if not self._log_condition_func(logs=logs, dataset_batch=None):
            return True

        self._init_logs(logs)

        return self._calc_whole_dataset_metrics(logs, f"{self._dataset_name}.dataset")

    def _init_logs(self, logs):
        current = self._get_logs_base(logs)
        if self._dataset_name not in current:
            current[self._dataset_name] = {}

        dataset_metrics = current[self._dataset_name]
        for metric_level in ['batch', 'window_average', 'dataset']:
            if metric_level not in dataset_metrics:
                dataset_metrics[metric_level] = {}

    def _get_current_evaluate_settings(self, logs):
        # Use custom settings if available, else use default settings

        evaluate_settings = get_value_at('evaluate_settings', logs, warn_on_failure=False)
        if evaluate_settings is None:
            evaluate_settings = self._evaluate_settings

        return evaluate_settings

    def _calc_whole_dataset_metrics(self, logs, log_path):

        current = self._get_logs_base(logs)
        metrics_log = get_value_at(log_path, current)

        evaluate_settings = self._get_current_evaluate_settings(logs)

        return self._metric_evaluator.calc_dataset_metrics_for(self._dataset,
                                                               metrics_log,
                                                               evaluate_settings=evaluate_settings,
                                                               dataset_name=self._dataset_name)

    def _init_metric_windows(self, reset=False):
        if reset:
            self._metric_windows = {}

        # Do nothing if metric windows are already available
        self._metric_windows = self._metric_windows or {}

        return True

    def _calc_batch_metric_data_from(self, batch, batch_metrics, logs):
        evaluate_settings = self._get_current_evaluate_settings(logs)

        model_output = None
        if self._dataset is None:
            current = self._get_logs_base(logs)
            loss = get_value_at(f"{self._dataset_name}.batch.loss", current)
            auxiliary_results = get_value_at(f"{self._dataset_name}.batch.auxiliary_results",
                                             current,
                                             warn_on_failure=False)

            model_output = {
                'loss': loss,
                'auxiliary_results': auxiliary_results
            }

        return self._metric_evaluator.calc_batch_metrics_for(batch,
                                                             batch_metrics,
                                                             evaluate_settings=evaluate_settings,
                                                             model_output=model_output)

    def _update_metrics_windows_for(self, metric_paths, batch_metrics, base_path):
        for metric_path in metric_paths:
            metric_value = get_value_at(metric_path, batch_metrics)

            full_metric_path = f"{base_path}.{metric_path}"
            sliding_window = self._metric_windows[metric_path] if metric_path in self._metric_windows else None
            if sliding_window is None:
                self._log.debug(f"Creating sliding window for {full_metric_path}")

                sliding_window = self._sliding_window_factory(length=self._batch_averaging_window, name=full_metric_path)
                self._metric_windows[metric_path] = sliding_window

            sliding_window.slide(metric_value)

    def _reduce(self, batch_metric_data_lists, reduced_metrics_log):

        return self._metric_evaluator.reduce(batch_metric_data_lists,
                                             reduced_metrics_log,
                                             dataset_name=self._dataset_name)

    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,
            "metric_evaluator": self._metric_evaluator is not None,
            "dataset_name": self._dataset_name,
            "batch_level": self._batch_level,
            "logging_mode": self._logging_mode,
            "dataset": self._dataset is not None,
            "evaluate_settings": self._evaluate_settings is not None,
            "batch_averaging_window": self._batch_averaging_window,
            "log_condition_func": self._log_condition_func is not None,
            "sliding_window_factory": self._sliding_window_factory is not None
        }

    def _check_state(self, state):
        state_attributes = ['metric_windows']

        for attr in state_attributes:
            if not has_key(state, attr):
                self._log.error(f"Given state does not have a value for {attr}, state is invalid")
                return False

        return True

    def _validate(self):
        self._valid = True

        if not hasattr(self._dataset, '__iter__'):
            if MetricsLoggingMode.will_log_window_average_metrics(self._logging_mode):
                self._log.debug("Training batch evaluation results will be used to calculate and log metrics.")
            else:
                self._log.error(f"No valid dataset provided to calculated metrics on, "
                                f"the {self} will not function")
                self._valid = False

        if not MetricEvaluatorBase.is_valid(self._metric_evaluator):
            self._log.error(f"The given metric evaluator is not valid, the {self} will not function : "
                            f"{self._metric_evaluator}")
            self._valid = False

        if not type(self._batch_level) is bool:
            self._log.error(f"Batch level must be True or False, given value is: {self._batch_level}\n"
                            f"the {self} will not function")
            self._valid = False

        if self._logging_mode not in MetricsLoggingMode:
            self._log.error(f"Metric logging mode unknown ({self._logging_mode})\n"
                            f"the {self} will not function")
            self._valid = False

        if not self._batch_level and self._logging_mode is not MetricsLoggingMode.WHOLE_DATASET_METRICS:
            self._log.error(f"For epoch level metric logging, the metric logging can only be "
                            f"MetricsLoggingMode.WHOLE_DATASET_METRICS. value is :{self._logging_mode}\n"
                            f"the {self} will not function")
            self._valid = False

        if self._log_condition_func is not None and not callable(self._log_condition_func):
            self._log.error(f"The log condition function must be callable, "
                            f"metric logging mode ({self._logging_mode}, the {self} will not function")
            self._valid = False


class TrainingMetricsLogger(MetricsLoggerBase):
    """
    By default, gets already calculated loss and auxiliary results from logs

    TODO : needs testing
    TODO : what to do for batch_level=false? Allow usage of window average of metrics to be used for epoch level?
    """

    def __init__(self,
                 metric_evaluator,
                 batch_level=True,
                 logging_mode=MetricsLoggingMode.BATCH_AND_WINDOW_AVERAGE_METRICS,
                 **kwargs):

        # print(f"mlpug.TrainingMetricsLogger(kwargs={kwargs})")

        super().__init__(
            metric_evaluator=metric_evaluator,
            dataset_name='training',
            batch_level=batch_level,
            logging_mode=logging_mode,
            **kwargs)

    def on_training_start(self,
                          num_epochs,
                          num_batches_per_epoch,
                          start_epoch,
                          start_batch,
                          start_update_iter):

        if MetricsLoggingMode.will_log_window_average_metrics(self._logging_mode):
            if self._batch_averaging_window is None:
                self._log.debug(f"Set batch_averaging_window to number of batches per epoch : {num_batches_per_epoch}")
                self._batch_averaging_window = num_batches_per_epoch

            if self._batch_averaging_window == math.inf:
                self._log.error("The batch average window is infinite, unable to calculate window average. "
                                "Please set a finite batch_averaging_window during construction of this "
                                "TrainingMetricsLogger.")
                self._valid = False
                return False

        return super().on_training_start(num_epochs, num_batches_per_epoch, start_epoch, start_batch, start_update_iter)


class TestMetricsLogger(MetricsLoggerBase):
    """
    Child class must implement _evaluate_loss method

    TODO : We should implement a specialization for the most common case of logging the (window/dataset) loss average
    """

    def __init__(self,
                 dataset,
                 dataset_name,
                 metric_evaluator,
                 batch_level=True,
                 logging_mode=None,
                 **kwargs):

        if logging_mode is None:
            logging_mode = MetricsLoggingMode.BATCH_AND_WINDOW_AVERAGE_METRICS \
                if batch_level else MetricsLoggingMode.WHOLE_DATASET_METRICS

        super().__init__(dataset=dataset,
                         dataset_name=dataset_name,
                         metric_evaluator=metric_evaluator,
                         batch_level=batch_level,
                         logging_mode=logging_mode,
                         **kwargs)

        self._dataset_iterator = None

    def on_training_start(self,
                          num_epochs,
                          num_batches_per_epoch,
                          start_epoch,
                          start_batch,
                          start_update_iter):

        if MetricsLoggingMode.will_log_window_average_metrics(self._logging_mode):
            if self._batch_averaging_window is None:
                try:
                    self._batch_averaging_window = len(self._dataset)
                    if self._batch_averaging_window == math.inf:
                        self._log.error(f"The batch average window of the {self._dataset_name} data set is infinite, "
                                        f"unable to calculate window average. "
                                        f"Please set a finite batch_processing_window during construction of "
                                        f"this TestMetricsLogger.")
                        self._valid = False
                    else:
                        self._log.debug(f"Set batch averaging window to number of batches in "
                                        f"{self._dataset_name} data set : {self._batch_averaging_window}")
                except Exception as e:
                    _.log_exception(self._log, f"Unable to assess data set length to set the batch averaging window, "
                                               f"{self} will not function", e)
                    self._valid = False

        return super().on_training_start(num_epochs, num_batches_per_epoch, start_epoch, start_batch, start_update_iter)

    def on_epoch_start(self, logs):
        """

        :param logs:

        :return: success (True or False)
        """
        success = super().on_epoch_start(logs)

        if not success:
            self._log.error('A problem occurred, will not continue executing this hook')
            return success

        if MetricsLoggingMode.will_log_window_average_metrics(self._logging_mode) and \
                self._dataset_iterator is None:
            self._dataset_iterator = iter(self._dataset)

        return True

    def on_batch_training_completed(self, training_batch, logs):
        current = self._get_logs_base(logs)
        training_iter = current["global_iter"]

        if self._logging_mode is MetricsLoggingMode.WHOLE_DATASET_METRICS:
            return super().on_batch_training_completed(training_batch, logs)
        else:
            try:
                dataset_batch = next(self._dataset_iterator)
            except StopIteration:
                self._log.debug(f"[Iteration {training_iter}] Rewinding the {self._dataset_name} data set")
                self._dataset_iterator = iter(self._dataset)
                dataset_batch = next(self._dataset_iterator)

            return super().on_batch_training_completed(dataset_batch, logs)

    def _validate(self):
        if not hasattr(self._dataset, '__iter__'):
            self._log.error(f"The given dataset {str(self._dataset)} is not iterable, the {self} will not function")
            self._valid = False

        super()._validate()
