# Standard library imports.
import os
import sys

import joblib

# Related third party imports.
import numpy as np
import scipy.stats
import torch
import torch.utils.data
from scipy.special import softmax
from sklearn import metrics

# Local application/library specific imports.
import util

def calc_evaluation_scores(y_true, y_pred, decision_threshold):
    """
    Calculate evaluation scores (AUC, pAUC, precision, recall, and F1 score)
    """
    auc = metrics.roc_auc_score(y_true, y_pred)
    p_auc = metrics.roc_auc_score(y_true, y_pred, max_fpr=CONFIG["max_fpr"])

    (_, false_positive, false_negative, true_positive,) = metrics.confusion_matrix(
        y_true, [1 if x > decision_threshold else 0 for x in y_pred]
    ).ravel()

    prec = true_positive / np.maximum(
        true_positive + false_positive, sys.float_info.epsilon
    )
    recall = true_positive / np.maximum(
        true_positive + false_negative, sys.float_info.epsilon
    )
    f1_score = 2.0 * prec * recall / np.maximum(prec + recall, sys.float_info.epsilon)

    print("AUC : {:.6f}".format(auc))
    print("pAUC : {:.6f}".format(p_auc))
    print("precision : {:.6f}".format(prec))
    print("recall : {:.6f}".format(recall))
    print("F1 score : {:.6f}".format(f1_score))

    return auc, p_auc, prec, recall, f1_score

def calc_performance_section(performance, csv_lines):
    """
    Calculate model performance per section.
    """
    amean_performance = np.mean(np.array(performance, dtype=float), axis=0)
    csv_lines.append(["arithmetic mean", ""] + list(amean_performance))
    hmean_performance = scipy.stats.hmean(
        np.maximum(np.array(performance, dtype=float), sys.float_info.epsilon),
        axis=0,
    )
    csv_lines.append(["harmonic mean", ""] + list(hmean_performance))
    csv_lines.append([])

    return csv_lines