#MachineLearning.py
#"""
#Copyright (c) 2021, Hiroaki Ito, kitasato-university
#Released under the MIT license
#https://opensource.org/licenses/mit-license.php
#"""

#Evaluation  1:peptide peak, 0:noise peak 

#Import libraries
import os
import sys
import pandas as pd
import numpy as np
import pathlib as ptl
import openpyxl
import copy
import xgboost as xgb
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
from sklearn.svm import LinearSVC
from sklearn.ensemble import RandomForestClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.naive_bayes import GaussianNB
from sklearn.neural_network import MLPClassifier
from sklearn.model_selection import KFold
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import confusion_matrix
from sklearn.decomposition import PCA
import seaborn as sns
import matplotlib.pyplot as plt


#Function for generating export for MachineLearning calculations
def machine_info(folder_test,prms,chk_ml,chk_mmsc,chk_pca,
                 pca_num,test_sz,srs,ml_name,cnfs,sen,spe,trainG,trainB,testG,testB):
    srs.rename('Count',inplace=True)
    df=pd.DataFrame([prms,['SVM','RF','XGB','ANN','KNN','GNB'],chk_ml,chk_mmsc,chk_pca,['Element','TestSize'],
                    [pca_num,test_sz],['Training Data Good','Training Data Bad','Test Data Good','Test Data Bad'],
                    [trainG,trainB,testG,testB]],
                    index=['Params','MachineName','Machine','MinMaxScaler','PCA','SettingName','Setting',
                    'Data','Number of Data']).T
    resl=srs.value_counts().sort_index(ascending=True).reset_index()
    resl.rename(columns={'index':'Results'},inplace=True)
    df_sheet1=pd.concat([df,resl],axis=1,sort=False)
    
    labels_index=['Predict Good','Predict Bad']
    labels_column=['Actual Good','Actual Bad']
    df_sheet2=pd.DataFrame()
    df_sheet3=pd.DataFrame()
    for i in range(len(ml_name)):
        mtrix=cnfs[i].T
        df_ins=pd.DataFrame(mtrix,index=labels_index,columns=labels_column).reset_index()
        df_ins.rename(columns={'index':ml_name[i]},inplace=True)
        df_sheet2=pd.concat([df_sheet2,df_ins],axis=1,sort=False)
        
        df_senspe=pd.DataFrame([sen[i],spe[i]],columns={ml_name[i]})
        df_sheet3=pd.concat([df_sheet3,df_senspe],axis=1,sort=False)
        
    df_sheet3.rename(index={0:'Sensitivity',1:'Specificity'},inplace=True)
    
    join_path=ptl.Path(folder_test,'Info_MachineLearning.xlsx')
    
    with pd.ExcelWriter(join_path) as writer:
        df_sheet1.to_excel(writer,sheet_name='LearningInfo')
        df_sheet2.to_excel(writer,sheet_name='ConfusionMatrix')
        df_sheet3.to_excel(writer,sheet_name='SensitivitySpecificity')
    
#Function for confusion matrix
def heatmap_matrix(matrix_data,ax):
    labels_index=['Predict1','Predict0']
    labels_column=['Actual1','Actual0']
    matrix_data=matrix_data.T
    df_matrix=pd.DataFrame(matrix_data,index=labels_index,columns=labels_column)
    plt.figure(figsize=(8,12))
    sns.heatmap(df_matrix,cmap='Blues',annot=True,annot_kws={"size":20},ax=ax)

#Function for SVM
def svm_(x_train,y_train,x_test,y_test,test_x,kf):
    #---------Support Vector Machine-------#
    svm=LinearSVC(tol=1e-2,max_iter=20000)
    param_svm={'C':np.linspace(1,100,10)}
    #Cross-Validation＋GRID SEARCH
    gs_svm=GridSearchCV(svm,param_grid=param_svm,cv=kf,scoring='accuracy')
    gs_svm.fit(x_train,y_train)
    
    y_pred_svm=gs_svm.predict(x_test)
    yx_pred_svm=gs_svm.predict(test_x)
    
    cnfs_svm=confusion_matrix(y_test,y_pred_svm,labels=[1,0])
    sensitivity_svm=cnfs_svm[0,0]/(cnfs_svm[0,0]+cnfs_svm[0,1])
    specificity_svm=cnfs_svm[1,1]/(cnfs_svm[1,0]+cnfs_svm[1,1])
    
    return 'SVM',yx_pred_svm,sensitivity_svm,specificity_svm,cnfs_svm
    #--------------------------------------#

#Function for RF
def rf_(x_train,y_train,x_test,y_test,test_x,kf):
    #---------Random Forest Classifier-------#
    rf=RandomForestClassifier(random_state=42)
    param_rf={'max_depth':[3,4,5],'n_estimators':[35,50,65,80]}
    gs_rf=GridSearchCV(rf,param_grid=param_rf,cv=kf,scoring='accuracy')
    gs_rf.fit(x_train,y_train)
    y_pred_rf=gs_rf.predict(x_test)
    yx_pred_rf=gs_rf.predict(test_x)
    
    cnfs_rf=confusion_matrix(y_test,y_pred_rf,labels=[1,0])
    sensitivity_rf=cnfs_rf[0,0]/(cnfs_rf[0,0]+cnfs_rf[0,1])
    specificity_rf=cnfs_rf[1,1]/(cnfs_rf[1,0]+cnfs_rf[1,1])
    return 'RF',yx_pred_rf,sensitivity_rf,specificity_rf,cnfs_rf
    #--------------------------------------#

#Function for XGBoost
def xgb_(x_train,y_train,x_test,y_test,test_x,kf):
    #---------------Xgboost----------------#
    param_xgb={'max_depth':[3,4,5],'n_estimators':[30,40]}
    model_xgb=xgb.XGBClassifier(eta=0.1)
    gs_xgb=GridSearchCV(model_xgb,param_grid=param_xgb,cv=kf,scoring='accuracy')
    gs_xgb.fit(x_train,y_train,early_stopping_rounds=10
               ,eval_metric='logloss',eval_set=[(x_test,y_test)])
    y_pred_xgb=gs_xgb.predict(x_test)
    yx_pred_xgb=gs_xgb.predict(test_x)
    
    cnfs_xgb=confusion_matrix(y_test,y_pred_xgb,labels=[1,0])
    sensitivity_xgb=cnfs_xgb[0,0]/(cnfs_xgb[0,0]+cnfs_xgb[0,1])
    specificity_xgb=cnfs_xgb[1,1]/(cnfs_xgb[1,0]+cnfs_xgb[1,1])
    
    return 'XGB',yx_pred_xgb,sensitivity_xgb,specificity_xgb,cnfs_xgb
    #--------------------------------------#

#Function for ANN
def ann_(x_train,y_train,x_test,y_test,test_x,kf):
    #---------------Neural Network----------------#
    clf_net=MLPClassifier(activation='logistic',solver='lbfgs')
    param_net={'hidden_layer_sizes':[[4,],[8,]],
               'alpha':[0.0001,0.001,0.01,0.1],
               'max_iter':[250,500]}
    #Cross-validation＋GRID SEARCH
    gs_net=GridSearchCV(clf_net,param_grid=param_net,cv=kf,scoring='accuracy')
    gs_net.fit(x_train,y_train)
    y_pred_net=gs_net.predict(x_test)
    yx_pred_net=gs_net.predict(test_x)
    
    cnfs_net=confusion_matrix(y_test,y_pred_net,labels=[1,0])
    sensitivity_net=cnfs_net[0,0]/(cnfs_net[0,0]+cnfs_net[0,1])
    specificity_net=cnfs_net[1,1]/(cnfs_net[1,0]+cnfs_net[1,1])
    
    return 'ANN',yx_pred_net,sensitivity_net,specificity_net,cnfs_net
    #---------------------------------------------#

#Function for KNN
def knn_(x_train,y_train,x_test,y_test,test_x,kf):
    #-----------K-NeighborsClassifier-------------#
    knn=KNeighborsClassifier()
    param_knn={'n_neighbors':[4,5,6]}
    #Cross-Validation＋GRID SEARCH
    gs_knn=GridSearchCV(knn,param_grid=param_knn,cv=kf,scoring='accuracy')
    gs_knn.fit(x_train,y_train)
    y_pred_knn=gs_knn.predict(x_test)
    yx_pred_knn=gs_knn.predict(test_x)
    
    cnfs_knn=confusion_matrix(y_test,y_pred_knn,labels=[1,0])
    sensitivity_knn=cnfs_knn[0,0]/(cnfs_knn[0,0]+cnfs_knn[0,1])
    specificity_knn=cnfs_knn[1,1]/(cnfs_knn[1,0]+cnfs_knn[1,1])
    
    return 'KNN',yx_pred_knn,sensitivity_knn,specificity_knn,cnfs_knn
    #--------------------------------------------#

#For GNB 
def gnb_(x_train,y_train,x_test,y_test,test_x,kf):
    #-----------Gaussian Naive Bayes-------------#
    gnb=GaussianNB()
    gnb.fit(x_train,y_train)
    y_pred_gnb=gnb.predict(x_test)
    yx_pred_gnb=gnb.predict(test_x)
    
    cnfs_gnb=confusion_matrix(y_test,y_pred_gnb,labels=[1,0])
    sensitivity_gnb=cnfs_gnb[0,0]/(cnfs_gnb[0,0]+cnfs_gnb[0,1])
    specificity_gnb=cnfs_gnb[1,1]/(cnfs_gnb[1,0]+cnfs_gnb[1,1])
    
    return 'GNB',yx_pred_gnb,sensitivity_gnb,specificity_gnb,cnfs_gnb
    #--------------------------------------------#

if __name__=='__main__':
    #file_train=training dataset path
    #file_test=validation dataset path
    #folder_test=validation dataset folder path
    file_train="C:\\Users\\Develop\\DDA_Integrate_machine_train.xlsx"
    file_test="C:\\Users\\Develop\\machine_for_paper\\DDA_Integrate_machine.xlsx"
    folder_test="C:\\Users\\Develop\\machine_for_paper"
    #list of features for ML
    prms=['idotp','AverageMassError(ppm)','MonoIso SN'
            ,'Sum PickMiss','Sum jagscore','FWHM_std'
           ,'PearsonCorr WeightedMean','CoelutionScore'
            ,'CoelutionCount']
    #DO min-max normalization and PCA except RF,XGB
    chk_ml=[True,True,True,True,True,True]
    chk_mmsc=[True,False,False,True,True,True]
    chk_pca=[True,False,False,True,True,True]
    #Dimensions for PCA
    pca_num=4
    
    test_sz=0

    #Import Training file
    df_train=pd.read_excel(file_train)
    #Import Test file
    df_test=pd.read_excel(file_test)

    df_x=df_train.loc[:,prms]
    df_t=df_test.loc[:,prms]
    
    if df_x.isnull().any().sum()>0 or df_t.isnull().any().sum()>0:
        #return to 0 if missing value
        df_x=df_x.fillna(0)
        df_t=df_t.fillna(0)
    
    ndf_x=df_x.values
    test_x=df_t.values
    #Training Evaluation
    if 'Evaluation' in df_train.columns:
        ndf_y=df_train['Evaluation'].values
    elif 'evaluation' in df_train.columns:
        ndf_y=df_train['evaluation'].values
    elif 'EVALUATION' in df_train.columns:
        ndf_y=df_train['evaluation'].values
    else:
        raise Exception('No column of Evaluation is exist.')

    #Divide to traing and cross-validations
    if test_sz>0:
        x_train,x_test,y_train,y_test=train_test_split(ndf_x,ndf_y,test_size=test_sz,random_state=42)
        train_num_good=np.count_nonzero(y_train==1)
        train_num_bad=np.count_nonzero(y_train==0)
        test_num_good=np.count_nonzero(y_test==1)
        test_num_bad=np.count_nonzero(y_test==0)
    
    else:
        x_train=ndf_x
        x_test=ndf_x
        y_train=ndf_y
        y_test=ndf_y
        train_num_good=np.count_nonzero(y_train==1)
        train_num_bad=np.count_nonzero(y_train==0)
        test_num_good=train_num_good
        test_num_bad=train_num_bad

    #---------MinMax Scaler---------#
    mmsc=MinMaxScaler()
    mmsc.fit(x_train)
    x_train_mmsc=mmsc.transform(x_train)
    x_test_mmsc=mmsc.transform(x_test)
    test_mmsc=mmsc.transform(test_x)
    #-------------------------------#
   
    #------PCA------#
    pca=PCA(n_components=pca_num)
    x_train_pca=pca.fit_transform(x_train_mmsc)
    x_test_pca=pca.transform(x_test_mmsc)
    test_pca=pca.transform(test_mmsc)

    #-----PCA (w/o normalization) ----#
    x_train_pca_only=pca.fit_transform(x_train)
    x_test_pca_only=pca.transform(x_test)
    test_pca_only=pca.transform(test_x)
    #---------------------#
    
    kf=KFold(n_splits=5,shuffle=True,random_state=42)
    
    definic_list=[svm_,rf_,xgb_,ann_,knn_,gnb_]
    ml_name=[]
    pred_list=[]
    sensitivity_list=[]
    specificity_list=[]
    cnfs_list=[]
    
    #Supervised Machine Learning  
    for i,ele1,ele2,ele3 in zip(range(len(definic_list)),chk_ml,chk_mmsc,chk_pca):
        val_i=15*i 
        if ele1==False:
            continue
        else:
            if ele2==False:
                if ele3==False:
                    nm,pred,sen,spe,cnfs=definic_list[i](x_train,y_train,x_test,y_test,test_x,kf)
                    ml_name.append(nm)
                    pred_list.append(pred)
                    sensitivity_list.append(sen)
                    specificity_list.append(spe)
                    cnfs_list.append(cnfs)
                else:
                    nm,pred,sen,spe,cnfs=definic_list[i](x_train_pca_only,y_train,x_test_pca_only,y_test,test_pca_only,kf)
                    ml_name.append(nm)
                    pred_list.append(pred)
                    sensitivity_list.append(sen)
                    specificity_list.append(spe)
                    cnfs_list.append(cnfs)
            else:
                if ele3==False:
                    nm,pred,sen,spe,cnfs=definic_list[i](x_train_mmsc,y_train,x_test_mmsc,y_test,test_mmsc,kf)
                    ml_name.append(nm)
                    pred_list.append(pred)
                    sensitivity_list.append(sen)
                    specificity_list.append(spe)
                    cnfs_list.append(cnfs)
                else:
                    nm,pred,sen,spe,cnfs=definic_list[i](x_train_pca,y_train,x_test_pca,y_test,test_pca,kf)
                    ml_name.append(nm)
                    pred_list.append(pred)
                    sensitivity_list.append(sen)
                    specificity_list.append(spe)
                    cnfs_list.append(cnfs)
    
    #Create directory for output
    mk_folder=folder_test+'\\MachineLearningFigures'
    folder_cnt=1
    if os.path.exists(mk_folder)==False:
        os.makedirs(mk_folder,exist_ok=True)
    else:
        while(os.path.exists(mk_folder)==True):
            mk_folder=folder_test+'\\MachineLearningFigures('+str(folder_cnt)+')'
            folder_cnt=folder_cnt+1
        os.makedirs(mk_folder,exist_ok=True)

    #----Plotting(Save plot figure png format)----#
    #Plot Confusion Matrix
    fig_mat,axis=plt.subplots(len(pred_list),1,figsize=(8,12))
    fig_mat.tight_layout()
    if len(pred_list)>1:
        axis0=axis.ravel()
    #Export all results
    df_results=pd.DataFrame()
    for i in range(len(pred_list)):
        df_=pd.DataFrame(pred_list[i],columns={ml_name[i]})
        df_results=pd.concat([df_results,df_],axis=1)
        if len(pred_list)>1:
            axis0[i].set_title(ml_name[i])
            heatmap_matrix(cnfs_list[i],axis0[i])
        else:
            axis.set_title(ml_name[i])
            heatmap_matrix(cnfs_list[i],axis)

    df_results['Results']=df_results.sum(axis=1)/len(df_results.columns)
    
    df_tested=pd.concat([df_test,df_results],axis=1,sort=False)
    join_path3=ptl.Path(folder_test,'MachineLearning.xlsx')
    file_cnt=1
    if os.path.exists(join_path3)==False:
        df_tested.to_excel(join_path3)
    else:
        while(os.path.exists(join_path3)==True):
            filename_rep='MachineLearning('+str(file_cnt)+').xlsx'
            file_cnt=file_cnt+1
            join_path3=ptl.Path(folder_test,filename_rep)
        df_tested.to_excel(join_path3)

    #Plot Sensitivity、Specificity
    fig_s,ax0=plt.subplots(1,1,figsize=(7,7))
    ax0.grid(which='major',axis='both',color='gray',alpha=0.5,linestyle='--')
    ax0.plot(ml_name,sensitivity_list,color='royalblue',marker='o',label='Sensitivity')
    ax0.plot(ml_name,specificity_list,color='r',marker='o',label='Specificity')
    ax0.set_ylabel('Score')
    ax0.set_ylim([0,1])
    ax0.legend()
    #Output confusion matrix, Sensitivity, and Specificity
    fig_mat.savefig(os.path.join(mk_folder,'ConfusionMatrix.png'))
    fig_s.savefig(os.path.join(mk_folder,'Sensitivity_Specificity.png'))
    series_=df_results['Results']
    machine_info(mk_folder,prms,chk_ml,chk_mmsc,chk_pca,pca_num,test_sz
                 ,series_,ml_name,cnfs_list,sensitivity_list,specificity_list,
                 train_num_good,train_num_bad,test_num_good,test_num_bad)
    plt.close()