# -*- coding: utf-8 -*-
# This codes are migrated from Keras Tuner: https://keras-team.github.io/keras-tuner/.  
# The copyright belows to the Keras Tuner authors.


from __future__ import absolute_import, division, print_function, unicode_literals


import numpy as np
from tensorflow import keras



class MetricObservation(object):

    def __init__(self, value, step):
        if not isinstance(value, list):
            value = [value]
        self.value = value
        self.step = step

    def append(self, value):
        if not isinstance(value, list):
            value = [value]
        self.value += value

    def mean(self):
        return np.mean(self.value)

    def get_config(self):
        return {'value': self.value,
                'step': self.step}

    @classmethod
    def from_config(cls, config):
        return cls(**config)

    def __eq__(self, other):
        if not isinstance(other, MetricObservation):
            return False
        return (other.value == self.value and
                other.step == self.step)

    def __repr__(self):
        return 'MetricObservation(value={}, step={})'.format(
            self.value, self.step)


class MetricHistory(object):

    def __init__(self, direction='min'):
        if direction not in {'min', 'max'}:
            raise ValueError(
                '`direction` should be one of '
                '{"min", "max"}, but got: %s' % (direction,))
        self.direction = direction
        self._observations = {}

    def update(self, value, step):
        if step in self._observations:
            self._observations[step].append(value)
        else:
            self._observations[step] = MetricObservation(
                value, step=step)

    def get_best_value(self):
        values = list(
            obs.mean() for obs in self._observations.values())
        if not values:
            return None
        if self.direction == 'min':
            return np.nanmin(values)
        return np.nanmax(values)

    def get_best_step(self):
        best_value = self.get_best_value()
        if best_value is None:
            return None
        for obs in self._observations.values():
            if obs.mean() == best_value:
                return obs.step

    def get_history(self):
        return sorted(self._observations.values(),
                      key=lambda obs: obs.step)

    def set_history(self, observations):
        for obs in observations:
            self.update(obs.value, step=obs.step)

    def get_statistics(self):
        history = self.get_history()
        history_values = [obs.mean() for obs in history]
        if not len(history_values):
            return {}
        return {
            'min': float(np.nanmin(history_values)),
            'max': float(np.nanmax(history_values)),
            'mean': float(np.nanmean(history_values)),
            'median': float(np.nanmedian(history_values)),
            'var': float(np.nanvar(history_values)),
            'std': float(np.nanstd(history_values))
        }

    def get_last_value(self):
        history = self.get_history()
        if history:
            last_obs = history[-1]
            return last_obs.mean()
        else:
            return None

    def get_config(self):
        config = {'direction': self.direction,
                  'observations': [obs.get_config() for obs in self.get_history()]}
        return config

    @classmethod
    def from_config(cls, config):
        instance = cls(config['direction'])
        instance.set_history([MetricObservation.from_config(obs)
                              for obs in config['observations']])
        return instance


class MetricsTracker(object):

    def __init__(self, metrics=None):
        # str -> MetricHistory
        self.metrics = {}
        self.register_metrics(metrics)

    def exists(self, name):
        return name in self.metrics

    def register_metrics(self, metrics=None):
        metrics = metrics or []
        for metric in metrics:
            self.register(metric.name)

    def register(self, name, direction=None):
        if self.exists(name):
            raise ValueError('Metric already exists: %s' % (name,))
        if direction is None:
            direction = infer_metric_direction(name)
        self.metrics[name] = MetricHistory(direction)

    def update(self, name, value, step=0):
        value = float(value)
        if not self.exists(name):
            self.register(name)

        prev_best = self.metrics[name].get_best_value()
        self.metrics[name].update(value, step=step)
        new_best = self.metrics[name].get_best_value()

        improved = new_best != prev_best
        return improved

    def get_history(self, name):
        self._assert_exists(name)
        return self.metrics[name].get_history()

    def set_history(self, name, observations):
        assert type(observations) == list
        if not self.exists(name):
            self.register(name)
        self.metrics[name].set_history(observations)

    def get_best_value(self, name):
        self._assert_exists(name)
        return self.metrics[name].get_best_value()

    def get_best_step(self, name):
        self._assert_exists(name)
        return self.metrics[name].get_best_step()

    def get_statistics(self, name):
        self._assert_exists(name)
        return self.metrics[name].get_statistics()

    def get_last_value(self, name):
        self._assert_exists(name)
        return self.metrics[name].get_last_value()

    def get_direction(self, name):
        self._assert_exists(name)
        return self.metrics[name].direction

    def get_config(self):
        return {
            'metrics': {
                name: metric_history.get_config()
                for name, metric_history in self.metrics.items()}}

    @classmethod
    def from_config(cls, config):
        instance = cls()
        instance.metrics = {
            name: MetricHistory.from_config(metric_history)
            for name, metric_history in config['metrics'].items()}
        return instance

    def _assert_exists(self, name):
        if name not in self.metrics:
            raise ValueError('Unknown metric: %s' % (name,))


_MAX_METRICS = {
    'Accuracy', 'BinaryAccuracy',
    'CategoricalAccuracy', 'SparseCategoricalAccuracy',
    'TopKCategoricalAccuracy', 'SparseTopKCategoricalAccuracy',
    'TruePositives', 'TrueNegatives',
    'Precision', 'Recall', 'AUC',
    'SensitivityAtSpecificity', 'SpecificityAtSensitivity'
}

_MAX_METRIC_FNS = {
    'accuracy', 'categorical_accuracy', 'binary_accuracy',
    'sparse_categorical_accuracy'
}


def infer_metric_direction(metric):
    # Handle str input and get canonical object.
    if isinstance(metric, str):
        metric_name = metric
        if len(metric_name) > 4 and metric_name[:4] == 'val_':
            metric_name = metric_name[4:]
        if metric_name == 'loss':
            # Special-case the overall loss.
            return 'min'
        try:
            metric = keras.metrics.get(metric_name)
        except ValueError:
            # Default to minimization for unknown metric.
            return 'min'

    # Metric class or function.
    if isinstance(metric, keras.metrics.Metric):
        name = metric.__class__.__name__
        if name == 'MeanMetricWrapper':
            name = metric._fn.__name__
    else:
        name = metric.__name__

    if name in _MAX_METRICS or name in _MAX_METRIC_FNS:
        return 'max'
    return 'min'
