# import helper packages
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import corr_coeff_table as ct
import helpers as h
import bootstrap_p_val as p
from sklearn.metrics import roc_auc_score, roc_curve, auc

# BINARY MODEL COMPARISON FUNCTIONS
#---------------------------------#

# this function performs a t-test for classification
# model assessment as outlined in the paper
# What Predicts U.S. Recessions? by Weiling Liu and Emanuel Moench
# see pages 8-13 to get an overview of the model itself
# paper link: https://shorturl.at/clvZ9
def roc_t_stat(y_true, y_pred_1, y_pred_2, corr_method, roc_auc_fun=roc_auc_score):
    # call helper function to make initial checks
    h.check_passed_data(y_true, y_pred_1, y_pred_2)

    # INITIAL CHECKS COMPLETE
    #########################
    # dependency link: https://shorturl.at/hwxy6
    
    # get roc scores
    roc_auc_1 = roc_auc_fun(y_true, y_pred_1)
    roc_auc_2 = roc_auc_fun(y_true, y_pred_2)
    
    # get the q stats for each model
    model_1_q_1, model_1_q_2 = h.q_calculations(roc_auc_1)
    model_2_q_1, model_2_q_2 = h.q_calculations(roc_auc_2)
    
    # get the r value
    # see corr_coeff_table.py for implementation
    avg_corr = h.correlation_coef(y_true, y_pred_1, y_pred_2, corr_method)
    avg_area = (roc_auc_1 + roc_auc_2) / 2
    r = ct.find_r_val(avg_corr, avg_area)
    
    # get the true and false counts
    tp = np.sum(y_true == 1)
    tn = np.sum(y_true == 0)
    
    # get the variance for each model
    model_1_var = h.get_variance_t_stat(roc_auc_1, model_1_q_1, model_1_q_2, tp, tn)
    model_2_var = h.get_variance_t_stat(roc_auc_2, model_2_q_1, model_2_q_2, tp, tn)
    
    # calculate the t val
    t_stat = h.get_test_stat(roc_auc_1, roc_auc_2, model_1_var, model_2_var, r)
    
    return t_stat

# z-score function based on the
# groundbreaking work from Hanley and McNeil (1983)
def roc_z_score(y_true, y_pred_1, y_pred_2, corr_method, roc_auc_fun=roc_auc_score):
    # call helper function to make initial checks
    h.check_passed_data(y_true, y_pred_1, y_pred_2)

    # INITIAL CHECKS COMPLETE
    #########################
    # dependency link: https://shorturl.at/hwxy6
    
    # get roc scores
    roc_auc_1 = roc_auc_fun(y_true, y_pred_1)
    roc_auc_2 = roc_auc_fun(y_true, y_pred_2)
    
    # get the q stats for each model
    model_1_q_1, model_1_q_2 = h.q_calculations(roc_auc_1)
    model_2_q_1, model_2_q_2 = h.q_calculations(roc_auc_2)
    
    # get the r value
    # see corr_coeff_table.py for implementation
    avg_corr = h.correlation_coef(y_true, y_pred_1, y_pred_2, corr_method)
    avg_area = (roc_auc_1 + roc_auc_2) / 2
    r = ct.find_r_val(avg_corr, avg_area)
    
    # get the true and false counts
    tp = np.sum(y_true == 1)
    tn = np.sum(y_true == 0)
    
    # get the variance for each model
    model_1_var = h.get_variance_z_score(roc_auc_1, model_1_q_1, model_1_q_2, tp, tn)
    model_2_var = h.get_variance_z_score(roc_auc_2, model_2_q_1, model_2_q_2, tp, tn)
    
    # calculate the z-score
    z_score = h.get_test_stat(roc_auc_1, roc_auc_2, model_1_var, model_2_var, r)
    
    return z_score

# calculate p-value with bootstrapping
# method for comparing two models using 
# at a given significance level with bootstrapping
# H0: Model 1 performance is significantly different from Model 2
# H1: Model 1 is not significantly different from Model 2
# returns p-val and list of the differences between models
def boot_p_val(
    y_true,
    y_pred_1,
    y_pred_2,
    compare_fun=np.subtract,
    score_fun=roc_auc_score,
    sample_weight=None,
    n_resamples=5000,
    two_tailed=True,
    seed=None,
    reject_one_class_samples=True
    ):
    
    # call helper function to make initial checks
    h.check_passed_data(y_true, y_pred_1, y_pred_2)
    
    # INITIAL CHECKS COMPLETE
    #########################
    # dependency link: https://shorturl.at/hwxy6
    
    # call function
    return p.p_val(
           y_true=y_true,
           y_pred_1=y_pred_1,
           y_pred_2=y_pred_2,
           compare_fun=compare_fun,
           score_fun=score_fun,
           sample_weight=sample_weight,
           n_resamples=n_resamples,
           two_tailed=two_tailed,
           seed=seed,
           reject_one_class_samples=reject_one_class_samples
           )

# PLOT MULTIPLE AUROC GRAPHS
#--------------------------#

# returns plt object which can be further
# edited as the user wishes using the 
# matplotlib library
# docs: https://matplotlib.org/stable/index.html
def stacked_roc_plt(
        y_true,
        model_preds,
        model_names,
        roc_fun=roc_curve,
        auc_fun=auc,
        fig_size=(8,6),
        linewidth=2,
        linestyle='-',
        rand_guess_color='black'
        ):
    """
    Params
    ------
    y_true: arrary like list of binary values
    model_preds: array of list like objects holding model predictions
    model_names: array of strings containing names for the model predictions
                must be in the same order as the model_pred data
    fig_size: optional param for fig size
    roc_fun: func to calc roc_curve, must return fpr, tpr, and thresholds
    auc_fun: takes fpr and tpr as args to return area under curve
    linewidth: set line width for plots --> see Matplotlib docs for reference
    linestyle: set line style for plots --> see Matplotlib docs for reference
    rand_guess_color: color to set for random guess line --> see Matplotlib docs for reference
    
    Return val
    ---------
    func returns matplotlib plot object, does not add any styles other than
    fig size basic line elements
    """
    # make sure lengths match
    assert len(model_preds) == len(model_names), 'Length mismatch between models and model names'
    # initialize fig and loop over results
    # create plot
    plt.figure(figsize=fig_size)
    for model_y_pred, model_name in zip(model_preds, model_names):
        # get the fpr, trp, and thresholds
        fpr, tpr, _ = roc_fun(y_true, model_y_pred)
        # call method on fpr and tpr
        roc_auc = auc_fun(fpr, tpr)
        plt.plot(fpr, tpr, linewidth=linewidth, linestyle=linestyle, label=f'{model_name} ROC Curve (area = %0.2f)' % roc_auc)
        plt.plot([0, 1], [0, 1], color=rand_guess_color, linewidth=linewidth, linestyle=':')
    return plt

 
# THRESHOLD FUNCTIONS
#-------------------#

# get the optimal threshold for classification
# based on the roc curve
# works best for data that is NOT highly imbalanced
def optimal_threshold(y_true, y_pred):
    # calc the curve
    fpr, tpr, thresholds = roc_curve(y_true, y_pred)
    # find the optimal threshold
    best_idx = np.argmax(tpr - fpr)
    best_threshold = thresholds[best_idx]
    
    return best_threshold

# finding the optimal cutoff for highly imbalanced data
# can be more challenging, therefore the prodecure
# needs to be slightly adjusted
# original source: https://shorturl.at/ewMS3
# uses geometric mean method
def optimal_threshold_imbalanced(y_true, y_pred):
    # calc the curve
    fpr, tpr, thresholds = roc_curve(y_true, y_pred)
    # calculate means and get idx
    gmeans = np.sqrt(tpr * (1-fpr))
    gmeans = np.array(gmeans)
    best_idx = np.argmax(gmeans)
    best_threshold = thresholds[best_idx]
    
    return best_threshold

# NON-PARAMETRIC AUROC FUNCTION
#-----------------------------#

# non-parametric approach for
# estimating the AUCROC
# paper: Performance Evaluation of Zero Net-Investment Strategies (2010)
# by Oscar Jorda and Alan M. Taylor
# the original implementation of this 
# methodology was to compare two investment strategies
# paper link: https://shorturl.at/oEORV
def auroc_non_parametric(y_true, y_pred):
    # get true neg and pos counts
    n_1 = sum(1 for y in y_true if y == 1)
    n_0 = sum(1 for y in y_true if y == 0)
    # get geometric mean
    g_mean = 1 / (n_1 * n_0)
    # create df and seperate
    df = pd.DataFrame({'y_true': y_true, 'y_pred': y_pred})
    # get predictions at 0 and 1
    y_pred_true = df[df['y_true'] == 1]['y_pred']
    y_pred_false = df[df['y_true'] == 0]['y_pred']
    # find sum
    total_sum = 0
    for zi in y_pred_true:
        for xj in y_pred_false:
            if zi > xj:
                total_sum += 1
            elif zi == xj:
                total_sum += 0.5
    auroc = g_mean * total_sum
    
    return auroc