# -*- coding: utf-8 -*-
"""
Created on Fri May 28 13:53:46 2021

@author: Peter
"""
# -*- coding: utf-8 -*-
"""
Created on Thu Oct 29 21:12:19 2020

@author: Peter
"""
#                     PART1: nested-cv elasticnet
import os
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.model_selection import StratifiedShuffleSplit
from sklearn.linear_model import LogisticRegressionCV
from tqdm import tqdm
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score, roc_auc_score
import matplotlib.pyplot as plt
from sklearn.metrics import plot_roc_curve
from sklearn import metrics
from sklearn.feature_selection import SelectFromModel
import dill
import statistics
from sklearn.metrics import roc_curve
import scipy

szempont="roc_auc"
random_state=1
cv_outer = StratifiedShuffleSplit(n_splits=100, test_size=0.2, random_state=random_state)
cv=5
test_size=0.2
l1_ratios=[0.001, 0.005, 0.01, 0.05, 0.1, 0.2, 0.3, 0.5, 0.7, 0.9]

os.chdir("/data10/working_groups/balint_group/gargya.peter/R/uterus")

features = pd.read_table('uterus_rnaseq_VST.txt', sep="\t", index_col=0)
features = features[features['label']!="G2"]
features=features.replace("G1",0)
features=features.replace("G3",1)
print(features["label"].value_counts())

labels = np.array(features['label'])

features= features.drop('label', axis = 1)
X_train, X_test, y_train, y_test = train_test_split(features, labels, 
                                                    test_size = test_size, random_state = random_state)
X_train=pd.DataFrame(X_train)
X_test=pd.DataFrame(X_test)
print([X_train.shape,y_train.shape,X_test.shape,y_test.shape])

best_c=[]
best_l1=[]
max_auc=[]
valid_auc=[]
tprs = []
aucs = []
mean_fpr = np.linspace(0, 1, 100)
fig, ax = plt.subplots()
i=1
genes=[]

for train_idx, val_idx in tqdm(cv_outer.split(X_train, y_train)):
    train_data, val_data = X_train.iloc[train_idx], X_train.iloc[val_idx]
    train_target, val_target = y_train[train_idx], y_train[val_idx]
    
    best_model = LogisticRegressionCV(random_state=random_state, cv=cv, scoring=szempont, penalty="elasticnet", 
                          fit_intercept=True, solver="saga", Cs=10, n_jobs=15, max_iter=20000, 
                          l1_ratios=l1_ratios)
    best_model.fit(train_data, train_target)
    
    model = SelectFromModel(best_model, prefit=True, threshold=None)
    mask=model.get_support()
    train_data_genes = train_data.loc[:, mask]
    [genes.append(x) for x in train_data_genes.columns.values]
    coefs_array=+ best_model.coef_
    
    y_pred_prob = best_model.predict_proba(val_data)[:,1]
    valid_auc.append(roc_auc_score(val_target, y_pred_prob))
    best_c.append(best_model.C_)
    best_l1.append(best_model.l1_ratio_)
    max_auc.append(best_model.scores_[1].mean(axis=0).max())
    
    viz = plot_roc_curve(best_model, val_data, val_target,
                         name='ROC fold {}'.format(i), alpha=0.5, lw=1, ax=ax)
    interp_tpr = np.interp(mean_fpr, viz.fpr, viz.tpr)
    interp_tpr[0] = 0.0
    tprs.append(interp_tpr)
    aucs.append(viz.roc_auc)
    i=i+1


dill.dump_session('globalsave_part1.pkl')

ax.plot([0, 1], [0, 1], linestyle='--', lw=2, color='r', label='Chance', alpha=.8)
mean_tpr = np.mean(tprs, axis=0)
mean_tpr[-1] = 1.0
mean_auc = metrics.auc(mean_fpr, mean_tpr)
print("Mean AUC: ", mean_auc)
std_auc = np.std(aucs)
ax.plot(mean_fpr, mean_tpr, color='b', label=r'Mean ROC (AUC = %0.2f $\pm$ %0.2f)' % (mean_auc, std_auc), lw=2, alpha=.8)
std_tpr = np.std(tprs, axis=0)
tprs_upper = np.minimum(mean_tpr + std_tpr, 1)
tprs_lower = np.maximum(mean_tpr - std_tpr, 0)
ax.fill_between(mean_fpr, tprs_lower, tprs_upper, color='grey', alpha=.2, label=r'$\pm$ 1 std. dev.')
ax.set(xlim=[-0.05, 1.05], ylim=[-0.05, 1.05], title="Receiver operating characteristic")
plt.savefig("/data10/working_groups/balint_group/gargya.peter/R/uterus/crossval_roc.pdf")

fig, ax = plt.subplots()
ax.plot([0, 1], [0, 1], linestyle='--', lw=2, color='r', label='Chance', alpha=.8)
ax.plot(mean_fpr, mean_tpr, color='b', label=r'Mean ROC (AUC = %0.2f $\pm$ %0.2f)' % (mean_auc, std_auc), lw=2, alpha=.8)
ax.fill_between(mean_fpr, tprs_lower, tprs_upper, color='grey', alpha=.2, label=r'$\pm$ 1 std. dev.')
ax.set(xlim=[-0.05, 1.05], ylim=[-0.05, 1.05], title="Receiver operating characteristic")
ax.legend(loc="lower right")
plt.xlabel("False Positive Rate", fontsize=14)
plt.ylabel("True Positive Rate", fontsize=14)
ax.tick_params(axis='both', which='major', labelsize=14)
plt.savefig("/data10/working_groups/balint_group/gargya.peter/R/uterus/crossval_roc_mean_only.pdf")
plt.close('all')


res=pd.DataFrame({"c":best_c, "l1":best_l1, "max auc":max_auc, "valid auc":valid_auc})
print(res)
print(statistics.mean(max_auc))
print(statistics.mean(valid_auc))

#%%
#               PART2: retrain model on whole dataset, predict G2
import os
os.chdir("/data10/working_groups/balint_group/gargya.peter/R/uterus")

import dill                          
dill.load_session('globalsave_part1.pkl')


features = pd.read_table('uterus_rnaseq_VST.txt', sep="\t", index_col=0)
features = features[features['label']!="G2"]
features=features.replace("G1",0)
features=features.replace("G3",1)
print(features["label"].value_counts())

labels = np.array(features['label'])
features= features.drop('label', axis = 1)

X_train, X_test, y_train, y_test = train_test_split(features, labels,
                                                    test_size = test_size, random_state = random_state)
X_train=pd.DataFrame(X_train)
X_test=pd.DataFrame(X_test)
print([X_train.shape,y_train.shape,X_test.shape,y_test.shape])

best_model = LogisticRegressionCV(random_state=random_state, cv=cv, scoring=szempont, penalty="elasticnet", 
                          fit_intercept=True, solver="saga", Cs=10, n_jobs=15, max_iter=20000, 
                          l1_ratios=l1_ratios)
best_model.fit(X_train,y_train)
print("Best C: ", best_model.C_)
print("Best l1_ratio: ", best_model.l1_ratio_)
print ('Max auc_roc:', best_model.scores_[1].mean(axis=0).max())

dill.dump_session('globalsave_part2.pkl')

y_pred = best_model.predict(X_train)
y_pred_proba=best_model.predict_proba(X_train)[:,1]
fpr, tpr, thresholds = roc_curve(y_train, y_pred_proba, drop_intermediate=True)
optimal_idx = np.argmax(tpr - fpr)
optimal_threshold = thresholds[optimal_idx]
print(confusion_matrix(y_train,y_pred))
print(classification_report(y_train,y_pred))
print("Accuracy train: ", accuracy_score(y_train, y_pred))
print("AUC train: ", roc_auc_score(y_train, y_pred_proba))
print("Threshold value is:", optimal_threshold)
for i in range(len(y_pred)):
    if y_pred_proba[i]<optimal_threshold:
        y_pred[i]=0
    else: y_pred[i]=1

print(confusion_matrix(y_train,y_pred))
print(classification_report(y_train,y_pred))
print("Accuracy train after thresholding: ", accuracy_score(y_train, y_pred))


y_pred = best_model.predict(X_test)
y_pred_proba=best_model.predict_proba(X_test)[:,1]
print(confusion_matrix(y_test,y_pred))
print(classification_report(y_test,y_pred))
print("Accuracy test: ", accuracy_score(y_test, y_pred))
print("AUC test: ", roc_auc_score(y_test, y_pred_proba))
for i in range(len(y_pred)):
    if y_pred_proba[i]<optimal_threshold:
        y_pred[i]=0
    else: y_pred[i]=1

print(confusion_matrix(y_test,y_pred))
print(classification_report(y_test,y_pred))
print("Accuracy test after thresholding: ", accuracy_score(y_test, y_pred))


y_pred_proba=best_model.predict_proba(X_test)[:,1]
fpr, tpr, threshold = metrics.roc_curve(y_test, y_pred_proba)
roc_auc = metrics.auc(fpr, tpr)

fig, ax = plt.subplots()
ax.plot([0, 1], [0, 1], linestyle='--', lw=2, color='r', label='Chance', alpha=.8)
ax.plot(fpr, tpr, 'b', label = 'AUC = %0.2f' % roc_auc)
ax.set(xlim=[-0.05, 1.05], ylim=[-0.05, 1.05], title="Receiver operating characteristic")
ax.legend(loc = 'lower right')
plt.xlabel("False Positive Rate", fontsize=14)
plt.ylabel("True Positive Rate", fontsize=14)
ax.tick_params(axis='both', which='major', labelsize=14)
plt.savefig("/data10/working_groups/balint_group/gargya.peter/R/uterus/test_roc.pdf")
plt.close('all')



features = pd.read_table('uterus_rnaseq_VST_G2.txt', sep="\t", index_col=0)
print(features["label"].value_counts())
y_test = np.array(features['label'])
X_test= features.drop('label', axis = 1)

y_pred_proba=pd.DataFrame(best_model.predict_proba(X_test)[:,1])
y_pred_proba.columns=["pred_proba"]
y_pred_proba["samples"]=X_test.index.values
y_pred_proba.to_csv("/data10/working_groups/balint_group/gargya.peter/R/uterus/G2_preds.txt", sep='\t',
                    index=False)

#%%
#                 PART3: important overlapping genes during CV-s
import os
os.chdir("/data10/working_groups/balint_group/gargya.peter/R/uterus")

import dill                          
dill.load_session('globalsave_part1.pkl')

big_model_AUCs=valid_auc
coefs_array=coefs_array/100
total_coefs=np.absolute(coefs_array).sum()
coefs_pd=pd.DataFrame(data=np.absolute(coefs_array), columns=X_train.columns.values).transpose().sort_values(0, axis=0, ascending=False)

toplot=pd.DataFrame(data=coefs_array, columns=X_train.columns.values).transpose()
toplot.columns=["coefs"]
toplot["abs"]=np.absolute(coefs_array).transpose()
toplot=toplot.sort_values("abs", axis=0, ascending=False).head(12)
toplot.to_csv("to_plot_barchart.txt", index=True, sep="\t")

ori_X_train=pd.DataFrame(X_train)
ori_X_test=pd.DataFrame(X_test)
loop_auc=[]
loop_auc_test=[]
loop_pvals=[]
loop_cum_coef=[]

for i in range(1,200):
    valid_auc=[]
    valid_auc_test=[]
    gene_mask=coefs_pd.index.values[:i]
    X_train=ori_X_train[gene_mask]
    X_test=ori_X_test[gene_mask]
                                
    for train_idx, val_idx in tqdm(cv_outer.split(X_train, y_train)):
        train_data, val_data = X_train.iloc[train_idx], X_train.iloc[val_idx]
        train_target, val_target = y_train[train_idx], y_train[val_idx]
        
        best_model = LogisticRegressionCV(random_state=random_state, cv=cv, scoring=szempont, penalty="elasticnet",
                          fit_intercept=True, solver="saga", Cs=10, n_jobs=15, max_iter=20000, 
                          l1_ratios=l1_ratios)
        best_model.fit(train_data, train_target)
        
        y_pred_prob = best_model.predict_proba(val_data)[:,1]
        auc = roc_auc_score(val_target, y_pred_prob)
        valid_auc.append(auc)
        mean_auc=statistics.mean(valid_auc)
        
        best_model.fit(X_train, y_train)
        y_pred_prob = best_model.predict_proba(X_test)[:,1]
        auc_test = roc_auc_score(y_test, y_pred_prob)
        valid_auc_test.append(auc_test)
        mean_auc_test=statistics.mean(valid_auc_test)
        
    loop_auc.append(mean_auc)
    loop_auc_test.append(mean_auc_test)
    loop_pvals.append(scipy.stats.wilcoxon(big_model_AUCs, valid_auc, alternative="two-sided").pvalue)
    loop_cum_coef.append(coefs_pd[0][0:i].sum())


dill.dump_session('globalsave_part3.pkl')

data=pd.DataFrame(loop_auc, columns=["auc"])
data["auc_test"]=loop_auc_test
data["pval"]=loop_pvals
data["cum_coef"]=loop_cum_coef
data.to_csv("results.txt", index=False, sep="\t")

fig, ax = plt.subplots()
ax.plot(range(1,200), loop_auc, "brown", label="Mean AUCs")
ax.plot(range(1,200), loop_auc_test, "green", label="Test AUCs")
plt.ylabel('AUC scores', fontsize=14)
plt.xlabel('Number of genes', fontsize=14)
plt.legend()
plt.title("Results of iterative analysis")
ax.tick_params(axis='both', which='major', labelsize=14)
plt.savefig("/data10/working_groups/balint_group/gargya.peter/R/uterus/loop_AUCs_min_num_genes.pdf")
plt.close('all')



#%%
#                   PART4: results with min num genes
import os
os.chdir("/data10/working_groups/balint_group/gargya.peter/R/uterus")

import dill                          
dill.load_session('globalsave_part1.pkl')

big_model_AUCs=valid_auc
coefs_pd=pd.DataFrame(data=np.absolute(coefs_array), columns=X_train.columns.values).transpose().sort_values(0, axis=0, ascending=False)

gene_mask=coefs_pd.index.values[:12]
print(gene_mask)
X_train=X_train[gene_mask]
X_test=X_test[gene_mask]

best_model = LogisticRegressionCV(random_state=random_state, cv=cv, scoring=szempont, penalty="elasticnet",
                                  fit_intercept=True, solver="saga", Cs=10, n_jobs=15, max_iter=20000, 
                                  l1_ratios=l1_ratios)
best_model.fit(X_train, y_train)
print("Best C: ", best_model.C_)
print("Best l1_ratio: ", best_model.l1_ratio_)
print ('Max auc_roc:', best_model.scores_[1].mean(axis=0).max())


y_pred = best_model.predict(X_train)
y_pred_proba=best_model.predict_proba(X_train)[:,1]
fpr, tpr, thresholds = roc_curve(y_train, y_pred_proba, drop_intermediate=True)
optimal_idx = np.argmax(tpr - fpr)
optimal_threshold = thresholds[optimal_idx]
print(confusion_matrix(y_train,y_pred))
print(classification_report(y_train,y_pred))
print("Accuracy train: ", accuracy_score(y_train, y_pred))
print("AUC train: ", roc_auc_score(y_train, y_pred_proba))
print("Threshold value is:", optimal_threshold)
for i in range(len(y_pred)):
    if y_pred_proba[i]<optimal_threshold:
        y_pred[i]=0
    else: y_pred[i]=1

print(confusion_matrix(y_train,y_pred))
print(classification_report(y_train,y_pred))
print("Accuracy train after thresholding: ", accuracy_score(y_train, y_pred))


train_pred_proba=pd.DataFrame(y_pred_proba)
train_pred_proba["samples"]=X_train.index.values


y_pred = best_model.predict(X_test)
y_pred_proba=best_model.predict_proba(X_test)[:,1]
print(confusion_matrix(y_test,y_pred))
print(classification_report(y_test,y_pred))
print("Accuracy test: ", accuracy_score(y_test, y_pred))
print("AUC test: ", roc_auc_score(y_test, y_pred_proba))
for i in range(len(y_pred)):
    if y_pred_proba[i]<optimal_threshold:
        y_pred[i]=0
    else: y_pred[i]=1

print(confusion_matrix(y_test,y_pred))
print(classification_report(y_test,y_pred))
print("Accuracy test after thresholding: ", accuracy_score(y_test, y_pred))


y_pred_proba=best_model.predict_proba(X_test)[:,1]
fpr, tpr, threshold = metrics.roc_curve(y_test, y_pred_proba)
roc_auc = metrics.auc(fpr, tpr)
fig, ax = plt.subplots()
ax.plot([0, 1], [0, 1], linestyle='--', lw=2, color='r', label='Chance', alpha=.8)
ax.plot(fpr, tpr, 'b', label = 'AUC = %0.2f' % roc_auc)
ax.set(xlim=[-0.05, 1.05], ylim=[-0.05, 1.05], title="Receiver operating characteristic")
ax.legend(loc = 'lower right')
plt.xlabel("False Positive Rate", fontsize=14)
plt.ylabel("True Positive Rate", fontsize=14)
ax.tick_params(axis='both', which='major', labelsize=14)
plt.savefig("/data10/working_groups/balint_group/gargya.peter/R/uterus/test_roc_with_min_num_genes.pdf")
plt.close('all')


features = pd.read_table('uterus_rnaseq_VST_G2.txt', sep="\t", index_col=0)
print(features["label"].value_counts())
y_test = np.array(features['label'])
X_test= features.drop('label', axis = 1)

X_test=X_test[gene_mask]

y_pred_proba=pd.DataFrame(best_model.predict_proba(X_test)[:,1])
y_pred_proba.columns=["pred_proba"]
y_pred_proba["samples"]=X_test.index.values
y_pred_proba.to_csv("/data10/working_groups/balint_group/gargya.peter/R/uterus/G2_preds_with_mingenes.txt", sep='\t',
                    index=False)