import numpy as np
from matplotlib import pyplot as plt

prop_cycle = plt.rcParams['axes.prop_cycle']
COLOR_CYCLE = prop_cycle.by_key()['color']


def _smooth(values, std):
    """Smooths a list of values by convolving with a gussian.
    Assumes equal spacing.
    Args:
      values: A 1D array of values to smooth.
      std: The standard devistion of the gussian. The units are array elements.
    Returns:
      The smoothed array.
    """
    width = std * 4
    x = np.linspace(-width, width, min(2 * width + 1, len(values)))
    kernel = np.exp(-(x / 5) ** 2)

    values = np.array(values)
    weights = np.ones_like(values)

    smoothed_values = np.convolve(values, kernel, mode='same')
    smoothed_weights = np.convolve(weights, kernel, mode='same')

    return smoothed_values / smoothed_weights


class HistoryPlotter(object):
    """A class for plotting named set of keras-histories.
    The class maintains colors for each key from plot to plot.
    """

    def __init__(self, metric=None, smoothing_std=None):
        self.color_table = {}
        self.metric = metric
        self.smoothing_std = smoothing_std

    def plot(self, histories, metric=None, smoothing_std=None):
        """Plots a {name: history} dictionary of keras histories.
        Colors are assigned to the name-key, and maintained from call to call.
        Training metrics are shown as a solid line, validation metrics dashed.
        Args:
          histories: {name: history} dictionary of keras histories.
          metric: which metric to plot from all the histories.
          smoothing_std: the standard-deviaation of the smoothing kernel applied
            before plotting. The units are in array-indices.
        """
        if metric is None:
            metric = self.metric
        if smoothing_std is None:
            smoothing_std = self.smoothing_std

        for name, history in histories.items():
            # Remember name->color asociations.
            if name in self.color_table:
                color = self.color_table[name]
            else:
                color = COLOR_CYCLE[len(self.color_table) % len(COLOR_CYCLE)]
                self.color_table[name] = color

            train_value = history.history[metric]
            val_value = history.history['val_' + metric]
            if smoothing_std is not None:
                train_value = _smooth(train_value, std=smoothing_std)
                val_value = _smooth(val_value, std=smoothing_std)

            plt.plot(
                history.epoch,
                train_value,
                color=color,
                label=name.title() + ' Train')
            plt.plot(
                history.epoch,
                val_value,
                '--',
                label=name.title() + ' Val',
                color=color)

        plt.xlabel('Epochs')
        plt.ylabel(metric.replace('_', ' ').title())
        plt.legend()

        plt.xlim(
            [0, max([history.epoch[-1] for name, history in histories.items()])])
        plt.grid(True)
