# Based on auto_random_forest_antibiotic_20210421.py
# imbalanced-learn==0.7, scikit-learn==0.23.2

import sys
# sys.path.insert(0, '/media/gsm/KINGSTON/PycharmProjects/药物和化合物研究202012/自动机器学习02')


import matplotlib.pyplot as plt
import numpy as np
import os
import pandas as pd
import pickle
import skopt

from datetime import datetime
from imblearn.ensemble import BalancedRandomForestClassifier
from multiprocessing import cpu_count
from sklearn import metrics
from sklearn.model_selection import StratifiedKFold, RepeatedStratifiedKFold, cross_val_score


print("CPU core count：{}".format(cpu_count()))
n_cpu = cpu_count()-4


path1 = "/media/gsm/TempData/Data/A Deep Learning Approach to Antibiotic Discovery/整理数据/" # For molecular smiles strings
# path1 = "/media/gsm/TempData/Toy/Data/A Deep Learning Approach to Antibiotic Discovery/整理数据/"
path2 = "/media/gsm/TempData/Data/A Deep Learning Approach to Antibiotic Discovery/整理数据/Features/" # For molecular features
# path2 = "/media/gsm/TempData/Toy/Data/A Deep Learning Approach to Antibiotic Discovery/整理数据/Features/"
path3 = "/media/gsm/TempData/Model/A Deep Learning Approach to Antibiotic Discovery/Learn/" # For models
# path3 = "/media/gsm/TempData/Toy/Model/A Deep Learning Approach to Antibiotic Discovery/Learn/"


## Begin, read data and features for train
data_set0  = pd.read_csv(path1+"/S1B_1.csv")

RDK_finger_prints1 = pd.read_csv(path2+"RDK_finger_prints_S1B.csv")
MACCS_keys1 = pd.read_csv(path2+"MACCS_keys_S1B.csv")
# atom_pair1 = pd.read_csv(path2+"atom_pair_S1B.csv")
# topological_torsion1 = pd.read_csv(path2+"topological_torsion_S1B.csv")
avalon1 = pd.read_csv(path2+"avalon_S1B.csv")
ECFP4_1 = pd.read_csv(path2+"ECFP4_S1B.csv")
ECFP6_1 = pd.read_csv(path2+"ECFP6_S1B.csv")
TPAPF1 = pd.read_csv(path2+"TPAPF_keys_S1B.csv")
TPATF1 = pd.read_csv(path2+"TPATF_keys_S1B.csv")

finger_prints = {"RDK": RDK_finger_prints1,
                 "MACCS": MACCS_keys1,
                 # "atom_pair": atom_pair1,
                 # "topological_torsion": topological_torsion1,
                 "avalon": avalon1,
                 "ECPF4": ECFP4_1,
                 "ECPF6": ECFP6_1,
                 "TPAPF": TPAPF1,
                 "TPATF": TPATF1}
## End, read data and features for train


## Begin, make the set for learning
y_column_name = 'Activity'
finger_name = "RDK"

data_set1 = data_set0[["SMILES", y_column_name]]
finger_print = finger_prints[finger_name]

learn_set = pd.merge(data_set1, finger_print)
learn_set.dropna(subset = [y_column_name], inplace=True)
learn_set.reset_index()

X = learn_set.drop(["SMILES", y_column_name], 1)
Y = learn_set[y_column_name]
## End, make the set for learning


## Begin, parameters of model
parameters_dict = {"RDK": {"bootstrap": True, "class_weight": "balanced",
                           "criterion": "gini", "max_depth": 5627,
                           "min_samples_leaf": 1, "min_samples_split": 2,
                           "n_estimators": 3629}} # random_forest_antibiotic_20210511_RDK_b
parameters = parameters_dict["RDK"]
## End, parameters of model


## Begin, train with cross validation
# kfold = StratifiedKFold(n_splits=10, random_state=0, shuffle=True)
kfold = RepeatedStratifiedKFold(n_splits=5, n_repeats=10, random_state=0)
model = BalancedRandomForestClassifier(bootstrap=parameters["bootstrap"],
                                       class_weight=parameters["class_weight"],
                                       criterion=parameters["criterion"],
                                       max_depth=parameters["max_depth"],
                                       min_samples_leaf=parameters["min_samples_leaf"],
                                       n_estimators=parameters["n_estimators"],
                                       sampling_strategy='not majority',
                                       n_jobs=n_cpu, verbose=1)
scores = cross_val_score(model, X, Y, cv=kfold, scoring='roc_auc')
print("Cross validation roc-auc: ", scores.tolist())
print("Cross validation mean roc-auc: ", scores.mean())
## End, train with cross validation


## Begin, train with the data set divided in the specified way
# Begin, split set
nrow = X.shape[0]
i_rows0 = np.linspace(0, nrow-1, nrow, dtype="int32").tolist()
i_rows1 = [value for index, value in enumerate(i_rows0) if value%10 != 6] # 1, 5, 6, 8, 9
i_rows2 = [value for index, value in enumerate(i_rows0) if value%10 == 6]

X_train, Y_train = X.loc[i_rows1], Y[i_rows1]
X_test, Y_test = X.loc[i_rows2], Y[i_rows2]
# End, split set

# Begin, train and save model
model = BalancedRandomForestClassifier(bootstrap=parameters["bootstrap"],
                                       class_weight=parameters["class_weight"],
                                       criterion=parameters["criterion"],
                                       max_depth=parameters["max_depth"],
                                       min_samples_leaf=parameters["min_samples_leaf"],
                                       n_estimators=parameters["n_estimators"],
                                       sampling_strategy='not majority',
                                       n_jobs=n_cpu, verbose=1)
model.fit(X_train, Y_train)

model_name = path3+'random_forest_antibiotic_20210511c'+"_"+finger_name
pickle.dump(model, open(model_name, "wb")) # save model
model = pickle.load(open(model_name, "rb")) # load model
# End, train and save model

# Begin, test model
Y_predict = model.predict(X_test)
Y_proba_predict = model.predict_proba(X_test)

display = metrics.plot_roc_curve(model, X_test, Y_test)
print('type(display):',type(display))
plt.show()

print("confusion matrix:")
print(pd.DataFrame(metrics.confusion_matrix(Y_test, Y_predict)))
# End, test model
## End, train with the data set divided in the specified way



## Begin, read the data and features to be predicted
data_set0  = pd.read_csv(path1+"/S2B_1.csv")

RDK_finger_prints1 = pd.read_csv(path2+"RDK_finger_prints_S2B.csv")
MACCS_keys1 = pd.read_csv(path2+"MACCS_keys_S2B.csv")
# atom_pair1 = pd.read_csv(path2+"atom_pair_S2B.csv")
# topological_torsion1 = pd.read_csv(path2+"topological_torsion_S2B.csv")
avalon1 = pd.read_csv(path2+"avalon_S2B.csv")
ECFP4_1 = pd.read_csv(path2+"ECFP4_S2B.csv")
ECFP6_1 = pd.read_csv(path2+"ECFP6_S2B.csv")
TPAPF1 = pd.read_csv(path2+"TPAPF_keys_S2B.csv")
TPATF1 = pd.read_csv(path2+"TPATF_keys_S2B.csv")

finger_prints = {"RDK": RDK_finger_prints1,
                 "MACCS": MACCS_keys1,
                 # "atom_pair": atom_pair1,
                 # "topological_torsion": topological_torsion1,
                 "avalon": avalon1,
                 "ECPF4": ECFP4_1,
                 "ECPF6": ECFP6_1,
                 "TPAPF": TPAPF1,
                 "TPATF": TPATF1}
## End, read the data and features to be predicted


## Begin, make the set to be predicted
data_set1 = data_set0[["SMILES"]]
finger_print = finger_prints[finger_name]

predict_set = pd.merge(data_set1, finger_print)
predict_set.dropna(axis=0, how='any', inplace=True)
 # The moleculars will be removed, if their features including nan values.
predict_set.reset_index()
## End, make the set to be predicted


## Begin, predict
X = predict_set.drop(["SMILES"], 1)

Y_predict = model.predict(X)
Y_proba_predict = model.predict_proba(X)

proba_set = pd.DataFrame(Y_proba_predict)
proba_set.columns = ["Prob_0", "Prob_1"]
proba_set = pd.concat([predict_set["SMILES"], proba_set], axis=1)

csv_name = path3+'random_forest_antibiotic_20210511c'+"_Predict"+"_"+finger_name+".csv"
proba_set.to_csv(csv_name, index=False)
## End, predict