################################################################################
# Copyright (c) 2021 ContinualAI.                                              #
# Copyrights licensed under the MIT License.                                   #
# See the accompanying LICENSE file for terms.                                 #
#                                                                              #
# Date: 14-12-2020                                                             #
# Author(s): Lorenzo Pellegrini                                                #
# E-mail: contact@continualai.org                                              #
# Website: www.continualai.org                                                 #
################################################################################

from typing import Dict, Union, Iterable, Sequence, Tuple, TYPE_CHECKING, List

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.axes import Axes
from numpy import ndarray, arange
from torch import Tensor

if TYPE_CHECKING:
    from avalanche.training.strategies import BaseStrategy
    from avalanche.benchmarks.scenarios import Experience
    from avalanche.evaluation import PluginMetric


EVAL = "eval"
TRAIN = "train"


def default_cm_image_creator(confusion_matrix_tensor: Tensor,
                             display_labels: Sequence = None,
                             include_values=False,
                             xticks_rotation=0,
                             yticks_rotation=0,
                             values_format=None,
                             cmap='viridis',
                             image_title=''):
    """
    The default Confusion Matrix image creator.
    Code adapted from
    `Scikit learn <https://scikit-learn.org/stable/modules/generated/sklearn.metrics.plot_confusion_matrix.html>`_ # noqa

    :param confusion_matrix_tensor: The tensor describing the confusion matrix.
        This can be easily obtained through Scikit-learn `confusion_matrix`
        utility.
    :param display_labels: Target names used for plotting. By default, `labels`
        will be used if it is defined, otherwise the values will be inferred by
        the matrix tensor.
    :param include_values: Includes values in confusion matrix. Defaults to
        `False`.
    :param xticks_rotation: Rotation of xtick labels. Valid values are
        float point value. Defaults to 0.
    :param yticks_rotation: Rotation of ytick labels. Valid values are
        float point value. Defaults to 0.
    :param values_format: Format specification for values in confusion matrix.
        Defaults to `None`, which means that the format specification is
        'd' or '.2g', whichever is shorter.
    :param cmap: Must be a str or a Colormap recognized by matplotlib.
        Defaults to 'viridis'.
    :param image_title: The title of the image. Defaults to an empty string.
    :return: The Confusion Matrix as a PIL Image.
    """

    fig, ax = plt.subplots()

    cm = confusion_matrix_tensor.numpy()
    n_classes = cm.shape[0]
    im_ = ax.imshow(cm, interpolation='nearest', cmap=cmap)
    cmap_min, cmap_max = im_.cmap(0), im_.cmap(256)

    if include_values:
        text_ = np.empty_like(cm, dtype=object)

        # print text with appropriate color depending on background
        thresh = (cm.max() + cm.min()) / 2.0

        for i in range(n_classes):
            for j in range(n_classes):
                color = cmap_max if cm[i, j] < thresh else cmap_min

                if values_format is None:
                    text_cm = format(cm[i, j], '.2g')
                    if cm.dtype.kind != 'f':
                        text_d = format(cm[i, j], 'd')
                        if len(text_d) < len(text_cm):
                            text_cm = text_d
                else:
                    text_cm = format(cm[i, j], values_format)

                text_[i, j] = ax.text(
                    j, i, text_cm,
                    ha="center", va="center",
                    color=color)

    if display_labels is None:
        display_labels = np.arange(n_classes)

    fig.colorbar(im_, ax=ax)

    ax.set(xticks=np.arange(n_classes),
           yticks=np.arange(n_classes),
           xticklabels=display_labels,
           yticklabels=display_labels,
           ylabel="True label",
           xlabel="Predicted label")

    if image_title != '':
        ax.set_title(image_title)

    ax.set_ylim((n_classes - 0.5, -0.5))
    plt.setp(ax.get_xticklabels(), rotation=xticks_rotation)
    plt.setp(ax.get_yticklabels(), rotation=yticks_rotation)

    fig.tight_layout()
    return fig


SEABORN_COLORS = (
    (0.2980392156862745, 0.4470588235294118, 0.6901960784313725),
    (0.8666666666666667, 0.5176470588235295, 0.3215686274509804),
    (0.3333333333333333, 0.6588235294117647, 0.40784313725490196),
    (0.7686274509803922, 0.3058823529411765, 0.3215686274509804),
    (0.5058823529411764, 0.4470588235294118, 0.7019607843137254),
    (0.5764705882352941, 0.47058823529411764, 0.3764705882352941),
    (0.8549019607843137, 0.5450980392156862, 0.7647058823529411),
    (0.5490196078431373, 0.5490196078431373, 0.5490196078431373),
    (0.8, 0.7254901960784313, 0.4549019607843137),
    (0.39215686274509803, 0.7098039215686275, 0.803921568627451),
)


def repartition_pie_chart_image_creator(
    label2counts: Dict[int, List[int]],
    counters: List[int],
    colors: Union[ndarray, Iterable, int, float] = SEABORN_COLORS,
    fmt: str = "%1.1f%%",
):
    """
    Create a pie chart representing the labels repartition.

    :param label2counts: A dict holding the counts for each label, of the form
        {label: [count_at_step_0, count_at_step_1, ...]}. Only the last count of
        each label is used here.
    :param counters: (unused) The steps the counts were taken at.
    :param colors: The colors to use in the chart.
    :param fmt: Formatting used to display the text values in the chart.
    """
    fig, ax = plt.subplots()
    ax: Axes

    labels, counts = zip(*((label, c[-1]) for label, c in label2counts.items()))

    ax.pie(counts, labels=labels, autopct=fmt, colors=colors)

    fig.tight_layout()
    return fig


def repartition_bar_chart_image_creator(
    label2counts: Dict[int, List[int]],
    counters: List[int],
    colors: Union[ndarray, Iterable, int, float] = SEABORN_COLORS,
):
    """
    Create a bar chart representing the labels repartition.

    :param label2counts: A dict holding the counts for each label, of the form
        {label: [count_at_step_0, count_at_step_1, ...]}. Only the last count of
        each label is used here.
    :param counters: (unused) The steps the counts were taken at.
    :param colors: The colors to use in the chart.
    """
    fig, ax = plt.subplots()
    ax: Axes

    y = -arange(len(label2counts))
    labels, counts = zip(*((label, c[-1]) for label, c in label2counts.items()))
    total = sum(counts)

    ax.barh(y, width=counts, color=colors)
    ax.set_yticks(y)
    ax.set_yticklabels(labels)

    ax.set_xlabel("Number of exemplars")
    ax.set_ylabel("Class")

    for i, count in enumerate(counts):
        ax.text(count / 2, -i, f"{count/total:.1%}", va="center", ha="center")

    fig.tight_layout()
    return fig


def default_history_repartition_image_creator(
    label2counts: Dict[int, List[int]],
    counters: List[int],
    colors: Union[ndarray, Iterable, int, float] = SEABORN_COLORS,
):
    """
    Create a stack plot representing the labels repartition with their history.

    :param label2counts: A dict holding the counts for each label, of the form
        {label: [count_at_step_0, count_at_step_1, ...]}.
    :param counters: The steps the counts were taken at.
    :param colors: The colors to use in the chart.
    """
    fig, ax = plt.subplots()
    ax: Axes

    ax.stackplot(
        counters,
        label2counts.values(),
        labels=label2counts.keys(),
        colors=colors,
    )
    ax.legend(loc='upper left')
    ax.set_ylabel("Number of examples")
    ax.set_xlabel("step")

    fig.tight_layout()
    return fig


def stream_type(experience: 'Experience') -> str:
    """
    Returns the stream name from which the experience belongs to.
    e.g. the experience can be part of train or test stream.

    :param experience: the instance of the experience
    """

    return experience.origin_stream.name


def phase_and_task(strategy: 'BaseStrategy') -> Tuple[str, int]:
    """
    Returns the current phase name and the associated task label.

    The current task label depends on the phase. During the training
    phase, the task label is the one defined in the "train_task_label"
    field. On the contrary, during the eval phase the task label is the one
    defined in the "eval_task_label" field.

    :param strategy: The strategy instance to get the task label from.
    :return: The current phase name as either "Train" or "Task" and the
        associated task label.
    """

    task = strategy.experience.task_labels
    if len(task) > 1:
        task = None  # task labels per patterns
    else:
        task = task[0]

    if strategy.is_eval:
        return EVAL, task
    else:
        return TRAIN, task


def bytes2human(n):
    # http://code.activestate.com/recipes/578019
    # >>> bytes2human(10000)
    # '9.8K'
    # >>> bytes2human(100001221)
    # '95.4M'
    symbols = ('K', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y')
    prefix = {}
    for i, s in enumerate(symbols):
        prefix[s] = 1 << (i + 1) * 10
    for s in reversed(symbols):
        if n >= prefix[s]:
            value = float(n) / prefix[s]
            return '%.1f%s' % (value, s)
    return "%sB" % n


def get_metric_name(metric: 'PluginMetric',
                    strategy: 'BaseStrategy',
                    add_experience=False,
                    add_task=True):
    """
    Return the complete metric name used to report its current value.
    The name is composed by:
    metric string representation /phase type/stream type/task id
    where metric string representation is a synthetic string
    describing the metric, phase type describe if the user
    is training (train) or evaluating (eval), stream type describes
    the type of stream the current experience belongs to (e.g. train, test)
    and task id is the current task label.

    :param metric: the metric object for which return the complete name
    :param strategy: the current strategy object
    :param add_experience: if True, add eval_exp_id to the main metric name.
            Default to False.
    :param add_task: if True the main metric name will include the task
        information. If False, no task label will be displayed.
        If an int, that value will be used as task label for the metric name.
    """

    phase_name, task_label = phase_and_task(strategy)
    stream = stream_type(strategy.experience)
    base_name = '{}/{}_phase/{}_stream'.format(str(metric),
                                               phase_name, stream)
    exp_name = '/Exp{:03}'.format(strategy.experience.current_experience)

    if task_label is None and isinstance(add_task, bool):
        add_task = False
    else:
        if isinstance(add_task, bool) and add_task:
            task_name = '/Task{:03}'.format(task_label)
        elif isinstance(add_task, int):
            task_name = '/Task{:03}'.format(add_task)
            add_task = True

    if add_experience and not add_task:
        return base_name + exp_name
    elif add_experience and add_task:
        return base_name + task_name + exp_name
    elif not add_experience and not add_task:
        return base_name
    elif not add_experience and add_task:
        return base_name + task_name


__all__ = [
    "default_cm_image_creator",
    "phase_and_task",
    "get_metric_name",
    "stream_type",
    "bytes2human",
    "default_history_repartition_image_creator",
    "repartition_pie_chart_image_creator",
    "repartition_bar_chart_image_creator",
]
