from sklearn import metrics

import matplotlib.pyplot as plt

import seaborn as sns

def performance_evaluation_report(model, X_test, y_test, show_plot=False, labels=None, show_pr_curve=False):

    '''

    Function for creating a performance report of a classification model.

    Parameters

    ----------

    model : scikit-learn estimator

        A fitted estimator for classification problems.

    X_test : pd.DataFrame

        DataFrame with features matching y_test

    y_test : array/pd.Series

        Target of a classification problem.

    show_plot : bool

        Flag whether to show the plot

    labels : list

        List with the class names.

    show_pr_curve : bool

        Flag whether to also show the PR-curve. For this to take effect, 

        show_plot must be True.

    Return

    ------

    stats : pd.Series

        A series with the most important evaluation metrics

    '''

    y_pred = model.predict(X_test)

    y_pred_prob = model.predict_proba(X_test)[:, 1]

    cm = metrics.confusion_matrix(y_test, y_pred)

    tn, fp, fn, tp = cm.ravel()

    fpr, tpr, threshold = metrics.roc_curve(y_test, y_pred_prob)

    roc_auc = metrics.auc(fpr, tpr)

    precision, recall, thresholds = metrics.precision_recall_curve(

        y_test, y_pred_prob)

    pr_auc = metrics.auc(recall, precision)

    if show_plot:

        if labels is None:

            labels = ['Negative', 'Positive']

        N_SUBPLOTS = 3 if show_pr_curve else 2

        PLOT_WIDTH = 15 if show_pr_curve else 12

        PLOT_HEIGHT = 5 if show_pr_curve else 6

        fig, ax = plt.subplots(

            1, N_SUBPLOTS, figsize=(PLOT_WIDTH, PLOT_HEIGHT))

        fig.suptitle('Performance Evaluation', fontsize=16)

        sns.heatmap(cm, annot=True, fmt="d", linewidths=.5, cmap='BuGn_r', square=True, cbar=False, ax=ax[0],

                    annot_kws={"ha": 'center', "va": 'center'})

        ax[0].set(xlabel='Predicted label',

                  ylabel='Actual label', title='Confusion Matrix')

        ax[0].xaxis.set_ticklabels(labels)

        ax[0].yaxis.set_ticklabels(labels)

        ax[1].plot(fpr, tpr, 'b-', label=f'ROC-AUC = {roc_auc:.2f}')

        ax[1].set(xlabel='False Positive Rate',

                  ylabel='True Positive Rate', title='ROC Curve')

        ax[1].plot(fp/(fp+tn), tp/(tp+fn), 'ro',

                   markersize=8, label='Decision Point')

        ax[1].plot([0, 1], [0, 1], 'r--')

        ax[1].legend(loc='lower right')

        if show_pr_curve:

            ax[2].plot(recall, precision, label=f'PR-AUC = {pr_auc:.2f}')

            ax[2].set(xlabel='Recall', ylabel='Precision',

                      title='Precision-Recall Curve')

            ax[2].legend()

    stats = {'accuracy': metrics.accuracy_score(y_test, y_pred),

             'precision': metrics.precision_score(y_test, y_pred),

             'recall': metrics.recall_score(y_test, y_pred),

             'specificity': (tn / (tn + fp)),

             'f1_score': metrics.f1_score(y_test, y_pred),

             'cohens_kappa': metrics.cohen_kappa_score(y_test, y_pred),

             'roc_auc': roc_auc,

             'pr_auc': pr_auc}

    return stats