# 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
from xgboost import XGBClassifier


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, imbalanced random forest
# kfold = StratifiedKFold(n_splits=10, random_state=0, shuffle=True)
kfold = RepeatedStratifiedKFold(n_splits=5, n_repeats=10, random_state=0)
model = BalancedRandomForestClassifier() # Default hyperparameters
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()) # About 0.892
## End, train with cross validation, imbalanced random forest


## Begin, train with cross validation, XGBClassifier
# kfold = StratifiedKFold(n_splits=10, random_state=0, shuffle=True)
kfold = RepeatedStratifiedKFold(n_splits=5, n_repeats=10, random_state=0)
model = XGBClassifier() # Default hyperparameters
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()) # About 0.901
## End, train with cross validation, XGBClassifier