import numpy as np
import matplotlib.pyplot as plt
import itertools

__author__ = "Wei-Chung Huang"
__copyright__ = "Copyright 2018, The SJSU MSCS Master project"
__license__ = "MIT"
__version__ = "1.0.0"


def plot_confusion_matrix(cm, classes,
                          normalize=False,
                          title='Confusion matrix',
                          cmap=plt.cm.Blues):
    """
    This function prints and plots the confusion matrix.
    Normalization can be applied by setting `normalize=True`.
    """
    if normalize:
        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
        print("Normalized confusion matrix")
    else:
        print('Confusion matrix, without normalization')

    plt.figure(figsize=(10, 10))
    plt.imshow(cm, interpolation='nearest', cmap=cmap)
    plt.title(title)
    plt.colorbar()
    tick_marks = np.arange(len(classes))
    plt.xticks(tick_marks, classes, rotation=45)
    plt.yticks(tick_marks, classes)

    fmt = '.1f' if normalize else 'd'
    thresh = cm.max() / 2.
    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
        plt.text(j, i, format(cm[i, j], fmt),
                 horizontalalignment="center",
                 color="white" if cm[i, j] > thresh else "black")

    plt.tight_layout()
    plt.ylabel('Families')
    plt.xlabel('Predicted label')
    plt.show()


def plot_roc(x, y, multiple_curve=False, labels=None):
    """
    Draw a ROC curve
    :param x:  FPR
    :param y:  TPR
    :param multiple_curve: bool
        True when x and y are lists
    :param labels: list
        List of legends strings
    :return:
    """

    plt.title('ROC')
    if multiple_curve is False:
        roc_auc = np.trapz(y, x)
        plt.plot(x, y, 'b', label='AUC = %0.2f' % roc_auc)
    else:
        n = len(x)
        for i in range(n):
            roc_auc = np.trapz(y[i], x[i])
            roc_auc_str = 'AUC = %0.2f' % roc_auc
            label = roc_auc_str
            if labels is not None:
                label = labels[i]
            if i == n - 1:
                plt.plot(x[i], y[i], 'b', label=label)
            elif i == 0:
                plt.plot(x[i], y[i], 'g', label=label)
            else:
                plt.plot(x[i], y[i], '-.', label=label)
    plt.legend(loc='lower right')
    plt.plot([0, 1], [0, 1], 'r--')
    plt.xlim([0, 1])
    plt.ylim([0, 1])
    plt.ylabel('True Positive Rate')
    plt.xlabel('False Positive Rate')
    plt.show()


def plot_bar(auc, title, multiple_bar=True, labels=None, legends=None):
    """
    Plot barchart of AUC
    :param auc: list
        2D m-by-n list where m is the number of rate set and n is length of rates
    :param title:
    :param multiple_bar:
    :param labels:
    :param legends:
    :return:
    """
    if multiple_bar:
        N = len(auc[0])
        rects = list()
        width = 0.25  # the width of the bars
        ind = np.arange(N)  # the x locations for the groups
        fig, ax = plt.subplots()
        ax.yaxis.grid(zorder=0)
        for i in range(len(auc)):
            rects.append(ax.bar(ind + width * i, auc[i], width, zorder=3))

        # add some text for labels, title and axes ticks
        ax.set_ylabel('AUC')
        ax.set_ylim([0.4, 1.0])
        ax.set_title(title)
        ax.set_xticks(ind + width / 2)
        ax.set_xticklabels(labels, rotation=45)

        # for i in range(len(auc)):
        ax.legend(tuple(rects), tuple(legends), loc=4)

        fig.tight_layout()
        plt.show()


def plot_curve(y, xtick_names, xtick_rotation_deg, xlabel, ylabel, title, grid=True, show=True):
    """
    Simply plot or return a tight curve
    :param y: List of data
    :param xtick_names: Name of xticks (String)
    :param xtick_rotation_deg: Degree of rotation of xticks (Integer)
    :param xlabel: X-axis name
    :param ylabel: Y-axis name
    :param title: Title
    :param grid: Enable grid or not (optional, default = True)
    :param show: Flag to show the plot or return the figure (optional, default = True)
    :return:
    """

    plt.plot(y)
    plt.ylabel(ylabel)
    plt.xlabel(xlabel)
    plt.title(title)
    plt.xticks(np.arange(len(y)), xtick_names, rotation=xtick_rotation_deg)
    plt.xlim([0, len(y) - 1])
    if grid:
        plt.grid(linestyle='-', linewidth=1)
    plt.tight_layout()

    if show:
        plt.show()
    else:
        return plt.gcf()
