#!/usr/bin/env python
"""
The tools.stats module contains some useful tools for calculating useful statistics which aren't
easily available else where.
"""

def calculate_VIF(X):
    """
    A function to calculate variance inflation factors to
    investigate multicollinearity between predictor variables. 
    Function written by Osian Roberts.

    :param X: a 2D numpy array of shape = (nSamples, nPredictors).
    :returns: VIF_scores; A numpy array containing the VIF for each predictor variable.

    Interpretation of VIF scores (this is somewhat subjective):
    1 = No multicollinearity.
    1-5 = Moderate multicollinearity.
    > 5 = High multicollinearity.
    > 10 = This predictor should be removed from the model.
    
    Example::
    
    import numpy
    # create some fake predictors to test the function:
    x1 = numpy.random.uniform(1, 100, 500)  # this should be highly collinear with x3.
    x2 = numpy.random.uniform(1, 100, 500)  # this should be independent of x1.
    x3 = 0.85 * x1 + 50 # this should be highly collinear with x1.
    x4 = 0.85 * numpy.log10(x1) + 50 # this should be moderately collinear with x1.
    
    X = numpy.array([x1, x2, x3, x4]).T
    
    vif_scores = calculate_VIF(X)
        
    for i in range(vif_scores.size):
        print('VIF for x' + str(i + 1) + ':', vif_scores[i])
    
    """
    import numpy
    from sklearn.linear_model import LinearRegression
    
    print('Calculating VIF for {} predictors variables...'.format(X.shape[1]))
    # fit auxiliary regressions to each predictor in X:
    lm = LinearRegression()
    VIF_scores = []
    for i in range(X.shape[1]):
        # fit an OLS linear model to each predictor (y) using the remaining predictor variables (x):
        y = X[:, i]
        x = numpy.delete(X, i, axis=1)
        lm.fit(x, y)
        y_pred = lm.predict(x)
        del x

        # calculate the coefficient of determination:
        ss_tot = sum((y - numpy.mean(y))**2)
        ss_res = sum((y - y_pred)**2)
        r2 = 1 - (ss_res / ss_tot)
        del y, y_pred, ss_res, ss_tot

        # calculate VIF:
        if r2 == 1:
            r2 = 0.999  # this prevents division by zero errors when r2 == 1.
        VIF_scores.append(1 / (1 - r2))
        del r2
    del lm, X
    VIF_scores = numpy.array(VIF_scores).round(decimals=1)
    return VIF_scores


def bin_accuracy_scores_prob(y_true, y_prob):
    """
    A function to calculate accuracy measures for probabilistic responses with sklearn and scipy.
    Function written by Osian Roberts.

    :param y_true: binary class labels, where 0 is absence and 1 is presence.
    :param y_prob: probability of presence scores e.g., generated by a species distribution model.

    :returns: a list containing two arrays - metrics = names of test metrics. scores = test scores for each metric.

    Useful reference:
    https://machinelearningmastery.com/how-to-score-probability-predictions-in-python
    """
    import numpy
    # check inputs:
    if not isinstance(y_true, numpy.ndarray):
        y_true = numpy.array(y_true)
    if not isinstance(y_prob, numpy.ndarray):
        y_prob = numpy.array(y_prob)
    if y_true.ndim != 1:
        raise SystemExit('ERROR: the true labels are not in a 1D array.')
    if y_prob.ndim != 1:
        raise SystemExit('ERROR: the probability of presence values are not in a 1D array.')
    if y_true.size != y_prob.size:
        raise SystemExit('ERROR: unequal number of binary labels and probabilities.')

        # ensure that y_true contains binary labels (i.e. 0 or 1 values):
    y_true = y_true.astype('uint8')
    if numpy.min(y_true) != 0 or numpy.max(y_true) != 1:
        raise SystemExit('ERROR: the true labels are not binary (zero or one values).')

    from sklearn.metrics import roc_auc_score
    # calculates area under the receiver operating curve score.
    # A score of 0.5 shows the model is unable to discriminate between presence and absence.
    roc_auc = roc_auc_score(y_true, y_prob)

    from sklearn.metrics import average_precision_score
    # calculates area under the precision-recall curve. Perfect model = 1.0.
    average_precision = average_precision_score(y_true, y_prob)

    from sklearn.metrics import brier_score_loss
    # This is a quadratic loss function that calculates the mean squared error between
    # predicted probabilities and the true presence-absence (binary) labels.
    # A model with no false positives/negatives has a score of 0.0. Perfect model = 1.0.
    brier_score = brier_score_loss(y_true, y_prob)

    from sklearn.metrics import log_loss
    # The is logarithmic loss function that more heavily penalises false positives/negatives than the brier score.
    # A model with no false positives/negatives has a score of 0.0. There is no upper bound.
    log_loss_score = log_loss(y_true, y_prob)

    from scipy.stats import pointbiserialr
    # The point biserial correlation coefficient, range -1 to 1.
    # Quantifies the correlation between a binary and continuous variable.
    r = pointbiserialr(y_true, y_prob)[0]

    metrics = ['Test AUC', 'Point-Biserial r', 'Av. Precision', 'Brier Score', 'Log-Loss Score']
    scores = numpy.array([roc_auc, r, average_precision, brier_score, log_loss_score]).round(decimals=6)
    del roc_auc, r, average_precision, brier_score, log_loss_score, y_true, y_prob
    return [metrics, scores]


def accuracy_scores_binary(y_true, y_pred):
    '''
    A function to calculate accuracy measures for a binary classification.
    Function written by Osian Roberts.

    Parameters:
    :param y_true: observed binary labels, where 0 is absence and 1 is presence.
    :param y_pred: predicted binary labels, where 0 is absence and 1 is presence.

    :returns: a list containing two numpy.arrays - (metrics: name of test metrics, scores: test scores for each metric)

    Reference: See pages 253 - 255 in:
    Guisan et al. (2017). Habitat suitability and distribution models: with applications in R.
    '''
    import numpy

    # check inputs:
    if not isinstance(y_true, numpy.ndarray):
        y_true = numpy.array(y_true)
    if not isinstance(y_pred, numpy.ndarray):
        y_pred = numpy.array(y_pred)
    if y_true.ndim != 1:
        raise SystemExit('ERROR: the true labels are not in a 1D array.')
    if y_pred.ndim != 1:
        raise SystemExit('ERROR: the predicted labels are not in a 1D array.')
    if y_true.size != y_pred.size:
        raise SystemExit('ERROR: unequal number of binary labels.')

    # ensure that y_true, y_pred contain binary labels (i.e. 0 or 1 values):
    y_true = y_true.astype('uint8')
    y_pred = y_pred.astype('uint8')
    if numpy.min(y_true) != 0 or numpy.max(y_true) != 1:
        raise SystemExit('ERROR: the true labels are not binary (zero or one values).')
    if numpy.min(y_pred) != 0 or numpy.max(y_pred) != 1:
        raise SystemExit('ERROR: the predicted labels are not binary (zero or one values).')

    metrics = numpy.array(['Prevalence', 'Overall Diagnostic Power', 'Correct Classification Rate',
                           'Misclassification Rate', 'Presence Predictive Power', 'Absence Predictive Power',
                           'Accuracy', 'Balanced Accuracy', 'Sensitivity', 'Specificity', 'Precision', 'F1 Score',
                           'Matthews Correlation', 'Cohen Kappa', 'Normalised Mutual Information',
                           'Hanssen-Kuiper skill'])

    try:
        n_presence = numpy.where(y_true == 1)[0].size
        n_absence = numpy.where(y_true == 0)[0].size

        # calculate true-presence, true-absence, false-presence and false-absence:
        TP = numpy.where((y_true == 1) & (y_pred == 1))[0].size
        TA = numpy.where((y_true == 0) & (y_pred == 0))[0].size
        FP = numpy.where((y_true == 1) & (y_pred == 0))[0].size
        FA = numpy.where((y_true == 0) & (y_pred == 1))[0].size  # aka sweet FA!

        # proportion of presence records:
        prevalence = (TP / FA) / y_true.size

        # proportion of absence records:
        ODP = 1 - prevalence

        # correct classification & misclassification rate
        CCR = (TP + TA) / y_true.size
        MR = (FP + FA) / y_true.size

        # Sensitivity (aka Recall or True Positive Rate):
        sensitivity = TP / n_presence

        # false presence rate - inverse of sensitivity (redundant?)
        #FPR = 1  - sensitivity

        # Presence and absence predictive power:
        PPP = TP / (TP + FP)
        APP = TA / (TA + FA)

        # Specificity (aka True Negative Rate):
        specificity = TA / n_absence

        # false positive rate - inverse of specificity (redundant?)
        #FPR = 1 - specificity

        # Accuracy scores:
        accuracy = (TP + TA) / (n_presence + n_absence)
        balanced_accuracy = ((TP / n_presence) + (TA / n_absence)) / 2

        # precision:
        precision = TP / (TP + FP)

        # F1 score:
        f1_score = 2 * TP / ((2*TP) + FP + FA)

        # Matthews Correlation Coefficient:
        MCC = ((TP * TA) - (FP * FA)) / (((TP + FP) * (TP + FA) * (TA + FP) * (TA + FA))**0.5)

        # Hanssen-Kuiper skill (unreliable when TA is very large):
        TSS = sensitivity + specificity - 1
        del n_presence, n_absence, TP, TA, FP, FA

        from sklearn.metrics import normalized_mutual_info_score as nmi_score
        nmi_score = nmi_score(y_true, y_pred)

        # Cohen's Kappa (caution: sensitive to sample size and proportion of presence records):
        from sklearn.metrics import cohen_kappa_score as kappa
        kappa = kappa(y_true, y_pred)

        scores = numpy.array([prevalence, ODP, CCR, MR, PPP, APP, accuracy, balanced_accuracy, sensitivity,
                           specificity, precision, f1_score, MCC, kappa, nmi_score, TSS]).round(decimals=6)
        del prevalence, ODP, CCR, MR, PPP, APP, accuracy, balanced_accuracy, sensitivity
        del specificity, precision, f1_score, MCC, kappa, nmi_score, TSS
    except Exception:
        scores = numpy.zeros(len(metrics))

    if metrics.size == scores.size:
        return [metrics, scores]
    else:
        raise SystemExit('ERROR: unable to calculate accuracy metrics.')

