import os

import numpy as np

import matplotlib.pyplot as plt

from sklearn.metrics import (precision_recall_curve as pr_curve,

                             confusion_matrix,

                             precision_score,

                             accuracy_score,

                             recall_score,

                             roc_curve,

                             auc)

def roc_auc(ground_truth, inferences, experiment_dir):

    fpr, tpr, threshold = roc_curve(ground_truth, inferences)

    roc_auc = auc(fpr, tpr).item()

    plt.title('Receiver Operating Characteristic')

    plt.plot(fpr, tpr, 'b', label = 'AUC = %0.3f' % roc_auc)

    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.savefig(os.path.join(experiment_dir, 'roc_auc.png'))

    plt.clf()

    return round(roc_auc, 3)

def pr_auc(ground_truth, inferences, experiment_dir):

    precision, recall, threshold = pr_curve(ground_truth, inferences)

    #precision, recall, threshold = sorted(zip(precision, recall))

    pr_auc = auc(recall, precision).item()

    plt.title('Precision - Recall')

    plt.plot(recall, precision, 'b', label = 'AUC = %0.3f' % pr_auc)

    plt.legend(loc = 'lower right')

    plt.xlim([0, 1])

    plt.ylim([0, 1])

    plt.ylabel('Precision')

    plt.xlabel('Recall')

    plt.savefig(os.path.join(experiment_dir, 'pr_auc.png'))

    plt.clf()

    return round(pr_auc, 3)

def calc_metrics(ground_truth, inferences, normalize=None, threshold=0.5):

    y_pred = inferences > threshold

    #accuracy = accuracy_score(y_true=ground_truth, y_pred=y_pred).item()

    #recall = recall_score(y_true=ground_truth, y_pred=y_pred).item()

    #precision = precision_score(y_true=ground_truth, y_pred=y_pred).item()

    conf_mat = confusion_matrix(y_true=ground_truth,

                                y_pred=y_pred,

                                normalize=normalize)

    tn, fp, fn, tp = list(map(lambda x: x.item(), conf_mat.ravel()))

    tpr = tp / (tp + fn)

    tnr = tn / (tn + fp)

    ppv = tp / (tp + fp)

    npv = tn / (tn + fn)

    acc = (tp+tn) / (tp+tn+fp+fn)

    metrics = {'threshold': threshold,

               'TPR': tpr,

               'TNR': tnr,

               'PPV': ppv,

               'NPV': npv,

               'ACC': acc}

    return {k: round(v, 3) for k,v in metrics.items()}