# -*- coding: utf-8 -*-
"""
Created on Wed May  6  2023

@author: Guangchen Liu  email:gch_liu@163.com
"""
import time
starttime = time.time()
from sklearn.model_selection import GridSearchCV
import numpy as np
import pandas as pd
from pandas import DataFrame
import c1_data_preprocess
import seq2num
import scipy.io as scio
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler    
import joblib    
import xgboost as xgb  
from sklearn import metrics
from sklearn.metrics import accuracy_score,confusion_matrix,classification_report


#display all data
np.set_printoptions(threshold=np.inf)
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)

'''==============================================='''
'''Step1 input fasta data'''
#label 0  virus
##--------------------------------------------
#v01  dsDNA
data_path_01=r"data_train_virus.fasta"
input_data_01=c1_data_preprocess.start(data_path_01,"seqs")
input_num_01=len(input_data_01)

f1=open("preprossed_seq_name_v01.txt","w")
f1.writelines(input_data_01[0])
f1.close()
input_data_01[0].to_csv('preprossed_seq_name_v01.txt',index=False,header=0)
input_data_01[1].to_csv('preprossed_seq_data_v01.txt',index=False,header=0)

seq_test_01=[]
for line in open("preprossed_seq_data_v01.txt"):  # combine c0_data_preprocess.py output matrix
    seq_test_01.append(line)
num_seq_test_01=seq2num.seq2num(seq_test_01)

data_v01=np.nan_to_num(num_seq_test_01)

scio.savemat('data_v01.mat',{'data_v01':data_v01})
h_v01=pd.DataFrame(data_v01)
writer=pd.ExcelWriter('data_v01.xlsx')
h_v01.to_excel(writer)
writer.save()

virus = h_v01

# add label 0 to virus
n1 = len(virus)
zeros = [0 for i in range(n1)]
zeros_df1 = DataFrame(zeros)

zeros_df1.reset_index(drop=True, inplace=True)
virus.reset_index(drop=True, inplace=True)
virus_new = pd.concat([zeros_df1,virus],axis=1)

'''======================================================================'''
#label 1 human_noTE
#--------------------------------------------------
data_path_H_noTE=r"data_train_non_ERV.fasta"
input_data_H_noTE=c1_data_preprocess.start(data_path_H_noTE,"seqs")
input_num_H_noTE=len(input_data_H_noTE)
input_data_H_noTE.iloc[:,0].to_csv('preprossed_seq_name_H_noTE.txt',index=False,header=0)
input_data_H_noTE.iloc[:,1].to_csv('preprossed_seq_data_H_noTE.txt',index=False,header=0)

seq_test_H_noTE=[]
for line in open("preprossed_seq_data_H_noTE.txt"): 
    seq_test_H_noTE.append(line)
num_seq_test_H_noTE=seq2num.seq2num(seq_test_H_noTE)

data_H_noTE=np.nan_to_num(num_seq_test_H_noTE)

scio.savemat('data_H_noTE.mat',{'data_H_noTE':data_H_noTE})
h_H_noTE=pd.DataFrame(data_H_noTE)
writer=pd.ExcelWriter('data_H_noTE.xlsx')
h_H_noTE.to_excel(writer)
writer.save()

h_H_noTE = pd.read_csv('data_train_full_nonLTR_and_CCDS.txt',header=None)

n2 = len(h_H_noTE)
ones = [1 for i in range(n2)]
ones_df2 = DataFrame(ones)
ones_df2.reset_index(drop=True, inplace=True)
h_H_noTE.reset_index(drop=True, inplace=True)
human_noTE = pd.concat([ones_df2,h_H_noTE],axis=1)


'''======================================================================'''
#label 2 human_TE
data_path_H_TE=r"data_train_ERV.fasta"
input_data_H_TE=c1_data_preprocess.start(data_path_H_TE,"seqs")
input_num_H_TE=len(input_data_H_TE)
input_data_H_TE.iloc[:,0].to_csv('preprossed_seq_name_H_TE.txt',index=False,header=0)
input_data_H_TE.iloc[:,1].to_csv('preprossed_seq_data_H_TE.txt',index=False,header=0)

seq_test_H_TE=[]
for line in open("preprossed_seq_data_H_TE.txt"): 
    seq_test_H_TE.append(line)
num_seq_test_H_TE=seq2num.seq2num(seq_test_H_TE)

data_H_TE=np.nan_to_num(num_seq_test_H_TE)

scio.savemat('data_H_TE.mat',{'data_H_TE':data_H_TE})
h_H_TE=pd.DataFrame(data_H_TE)
writer=pd.ExcelWriter('data_H_TE.xlsx')
h_H_TE.to_excel(writer)
writer.save()

h_H_TE = pd.read_csv('data_LTR_train_total.txt',header=None)

n3 = len(h_H_TE)

twos = np.ones((n3,1))*2
twos_df2 = DataFrame(twos)
twos_df2.reset_index(drop=True, inplace=True)
h_H_TE.reset_index(drop=True, inplace=True)
human_TE = pd.concat([twos_df2,h_H_TE],axis=1)

'''==============================================='''
# merge all
virus_new.reset_index(drop=True, inplace=True)
h_H_noTE.reset_index(drop=True, inplace=True)
h_H_TE.reset_index(drop=True, inplace=True)
datatrain0 = pd.concat([virus_new,human_noTE,human_TE])
datatrain0.to_csv('datatrain_3labels_vNCBI_nonLTR_LTR_model_4_GPU.txt', index=False,header=None)

datatrain1 = pd.read_table('datatrain_3labels_vNCBI_nonLTR_LTR_model_4_GPU.txt', sep=',',header=None)
np.array(datatrain1)

datatrain1[np.isinf(datatrain1)] = np.nan
datatrain1.to_csv('datatrain_3labels_vNCBI_nonLTR_LTR_model_GPU.txt', index=False,header=None)

X = datatrain1.iloc[0:,1:]
y = datatrain1.iloc[0:,0]


'''==============================================='''
'''Step2 train and test data split'''
X_train,X_test,y_train,y_test=train_test_split(X,y,test_size=0.005,random_state=33)


'''==============================================='''
'''Step3 normalization'''    
scaler=MinMaxScaler().fit(X_train.astype(float))
 
joblib.dump(scaler,'scaler_3labels_vNCBI_nonLTR_LTR_model_4_GPU.m')   
X_train=scaler.transform(X_train)
X_test=scaler.transform(X_test)

'''==============================================='''
'''Step4 XGBoost parameters tuning and crossvalidation'''
model = xgb.XGBClassifier(tree_method='gpu_hist')#using gpu
param_dist = {"max_depth": list(range(4,10,2)),
              "min_child_weight" : list(range(0,6,2)),
              "n_estimators": [200],
              "learning_rate": [0.05, 0.1,0.2],
              "silent":[0],
              "subsample":[0.5,0.8],
              "colsample_bytree":[0.5,0.8],
           }

grid_search = GridSearchCV(model, param_grid=param_dist, cv = 3, 
                                   verbose=1)
grid_search.fit(X_train, y_train)

# collecting best parameters information
best_score = grid_search.best_score_
best_params = grid_search.best_params_
best_estimator = grid_search.best_estimator_
#  output results
print('\n best score is:\n',best_score,
      '\n\n best params are:\n',best_params,
      '\n\n best model:\n',best_estimator)

'''==============================================='''
'''Step5 model fitting'''
model_XGBoost = best_estimator
model_XGBoost.fit(X_train,y_train)  
     
joblib.dump(model_XGBoost,'model_xgboost_3labels_vNCBI_nonLTR_LTR_model_GPU.m')   
 
#### plot the features importance
#print('\n\n Feature importance:\n',model_XGBoost.feature_importances_)               #重要性排序
#from matplotlib import pyplot
#pyplot.barh(range(len(model_XGBoost.feature_importances_)), model_XGBoost.feature_importances_)
#pyplot.show()

'''==============================================='''
'''Step6 model evaluation'''
y_pred=best_estimator.predict(X_test)

### metrics by sklearn
accuracy = accuracy_score(y_test,y_pred)
confusion_matrix = confusion_matrix(y_test,y_pred,labels=[0,1,2])

precision = metrics.precision_score(y_test, y_pred,average='macro')
recall = metrics.recall_score(y_test, y_pred,average='macro')
f1_score = metrics.f1_score(y_test, y_pred,average='macro')

print('\n confusion_matrix = \n',confusion_matrix)
print('\n classification_report =\n', classification_report)
print('\n accuracy_sklearn = ',accuracy)
print('\n precison_sklearn = ',precision)
print('\n recall_sklearn = ',recall)
print('\n f1_Score_sklearn = ',f1_score)


 #save results
pred_result=pd.DataFrame({"real value":np.transpose(y_test)})    
evaluate_result=pd.DataFrame(columns=("Accuracy","Recall","Precision","F1-score","Time_consuming"))
# sort feature importance
feature_importance=pd.DataFrame({"feature_importance":np.transpose(model_XGBoost.feature_importances_ )})
fi_sort = feature_importance.sort_values(by='feature_importance',ascending = False) 

endtime = time.time()
testtime = endtime-starttime

pred_result['predict value']=np.transpose(y_pred)
evaluate_result.loc['XGBoost']=[
                     np.round(accuracy,3)
                    ,np.round(float(recall),3)
                    ,np.round(float(precision),3)
                    ,np.round(float(f1_score),3)                                                       
                    ,np.round(testtime,3)
                    ]

writer=pd.ExcelWriter("result_3labels_vNCBI_nonLTR_LTR_model_GPU.xlsx")
pred_result.to_excel(writer,sheet_name="predict")
evaluate_result.to_excel(writer,sheet_name="evaluate")
fi_sort.to_excel(writer,sheet_name="feature_importance_sort")
writer.save()

print("\n Time consuming is",np.round(testtime,3))
print('-----------------------------the end-----------------------------------')