import numpy as np
import warnings

from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import Lasso
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import roc_curve, auc, confusion_matrix

warnings.filterwarnings("ignore")


def make_confusion_matrix(prediction_probs, true_phenotypes, false_pos_rates,
                          true_pos_rates, thresholds):
    j_stats = {}
    for x, threshold in enumerate(thresholds):
        j_stat = true_pos_rates[x] + (1 - false_pos_rates[x]) - 1
        j_stats[j_stat] = threshold

    max_j_stat = max(j_stats)
    threshold = j_stats[max_j_stat]

    predictions = []
    for prob in prediction_probs:
        if prob >= threshold:
            predictions.append(1)
        else:
            predictions.append(0)

    matrix = confusion_matrix(y_true=true_phenotypes, y_pred=predictions)

    return [matrix, max_j_stat, threshold]


def kfold_cv(phenotypes, probes, probe_set):

    # phenotypes = list of subject phenotypes as [0, 1] based on classifier
    # probes =  list of probe names,
    # probe_set = 2d list of probe expression values for each subject

    lasso = Lasso(random_state=0, max_iter=10000)
    alphas = np.logspace(-4, -0.5, 30)

    tuned_parameters = [{'alpha': alphas}]
    n_folds = len(phenotypes) - 1  # subtracting 1 for outer cv

    prediction_probs = []
    coeff_dict = {}

    # pass scoring into gridsearch
    for leave_out in range(len(phenotypes)):
        probe_set_lo = probe_set[:leave_out] + probe_set[leave_out + 1:]
        phenotypes_lo = phenotypes[:leave_out] + phenotypes[leave_out + 1:]
        left_out = [probe_set[leave_out]]

        scaler = StandardScaler()
        scaler.fit(probe_set_lo)

        probe_set_lo = scaler.transform(probe_set_lo)
        probe_set_lo = probe_set_lo.tolist()
        left_out = scaler.transform(left_out)
        left_out = left_out.tolist()

        clf = GridSearchCV(lasso,
                           tuned_parameters,
                           cv=n_folds,
                           n_jobs=-1,
                           refit=True)
        clf.fit(probe_set_lo, phenotypes_lo)

        coeffs = clf.best_estimator_.coef_
        for x in range(len(coeffs)):
            if probes[x] not in coeff_dict:
                coeff_dict[probes[x]] = [coeffs[x]]
            else:
                coeff_dict[probes[x]].append(coeffs[x])

        prediction_val = clf.predict(left_out)
        probability = 1 / (1 + np.exp(-prediction_val))
        prediction_probs.append(probability)

    fpr, tpr, thresholds = roc_curve(phenotypes, prediction_probs)
    score = auc(fpr, tpr)
    confusion_matrix, j_stat, threshold = make_confusion_matrix(
        prediction_probs, phenotypes, fpr, tpr, thresholds)

    return [
        score, coeff_dict, confusion_matrix, j_stat, threshold,
        prediction_probs, phenotypes
    ]


def test_cross_validation():
    # Example for the Bacterial cross-validation of the HL2 signature (2 genes) on the GSE29161 datset.
    
    # subject_key = [GSM721702,GSM721705,GSM721706,GSM721708,GSM721710,GSM721712,GSM721713,GSM721714,GSM721715,GSM721716]
    subject_phenotypes = [1, 1, 1, 1, 1, 0, 0, 0, 0, 0]
    gene_names = ['IFI44L', 'FAM89A']
    gene_expression = [[10.653, 8.045], [8.236, 8.124], [8.622, 8.32],
                       [7.833, 7.842], [7.492, 8.971], [8.223, 8.289],
                       [8.099, 7.906], [7.491, 7.266], [7.101, 7.707],
                       [7.977, 7.32]]

    score, coeff_dict, confusion_matrix, j_stat, threshold, prediction_probs, phenotypes = kfold_cv(
        subject_phenotypes, gene_names, gene_expression)

    assert round(score,2) == 0.84, "Unexpected AUC Score"
    assert list(confusion_matrix[0]) == [3,2], "Unexpected Confusion Matrix"
    assert list(confusion_matrix[1]) == [0,5], "Unexpected Confusion Matrix"
    assert round(threshold, 3) == 0.567, "Unexpected J-stat Threshold"
    print("** All tests passed **")



test_cross_validation()