# Path: moloi/splits/scaffold_split.py
# def scaffold_split(smiles, frac_train=.8, seed=777):
#     """
#     Splits compounds into train/validation/test by scaffold.
# 
#     Warning: if there is one very popular scaffold can produce unbalanced split
# 
#     Params
#     ------
#     smiles: list
#         List of smiles
#     frac_train: float, default: 0.8
#         Float in [0, 1] range indicating size of the training set
#     seed: int
#         Used to shuffle smiles before splitting
# 
#     Notes
#     -----
#     Copied from https://github.com/deepchem/deepchem/blob/master/deepchem/splits/splitters.py
#     """
#     smiles = list(smiles)
# 
#     scaffolds = {}
# 
#     for ind, smi in tqdm(enumerate(smiles), total=len(smiles)):
#         scaffold = generate_scaffold(smi)
#         if scaffold not in scaffolds:
#             scaffolds[scaffold] = [ind]
#         else:
#             scaffolds[scaffold].append(ind)
# 
#     scaffolds_keys = list(scaffolds)
#     rng = np.random.RandomState(seed)
#     rng.shuffle(scaffolds_keys)
# 
#     train_cutoff = frac_train * len(smiles)
#     train_inds, test_inds = [], []
#     for scaffold_key in scaffolds_keys:
#         if len(train_inds) > train_cutoff:
#             test_inds += scaffolds[scaffold_key]
#         else:
#             train_inds += scaffolds[scaffold_key]
#     return train_inds, test_inds
# 
# Path: moloi/splits/cluster_split.py
# def cluster_split(smi, test_cluster_id, n_splits):
#     X_molprint, _ = molprint2d_count_fingerprinter(np.array(smi))
#     X_molprint = X_molprint.astype(np.float32)
#     K_molprint_tanimoto = tanimoto_similarity(X_molprint, X_molprint)
# 
#     clustering = BalancedAgglomerativeClustering(n_clusters=n_splits)
#     logger.info("Clustering")
#     labels = clustering.fit_predict(K_molprint_tanimoto)
# 
#     train_ids = np.where(np.array(labels) != test_cluster_id)[0]
#     test_ids = np.where(np.array(labels) == test_cluster_id)[0]
# 
#     return train_ids, test_ids
# 
# Path: moloi/config_processing.py
# def cv_splits_save(split_type, split_size, n_cv, data_config, targets):
#     sv_config = ConfigParser.ConfigParser()
# 
#     sv_config.read(data_config)
#     try:
#         sv_config[split_type + " " + str(split_size)]['cv_' + str(targets)] = str(n_cv)
#     except:
#         with open(data_config, "a") as ini:
#             ini.write('[' + split_type + " " + str(split_size) + ']')
#         sv_config.read(data_config)
#         sv_config[split_type + " " + str(split_size)]['cv_' + str(targets)] = str(n_cv)
#     with open(data_config, 'w') as configfile:
#         sv_config.write(configfile)
# 
# def cv_splits_load(split_type, split_size, data_config, targets):
#     loaded_cv = False
#     cv_config = ConfigParser.ConfigParser()
#     cv_config.read(data_config)
#     if split_type is False:
#         section = 'init'
#     else:
#         section = split_type + " " + str(split_size)
# 
#     if split_type + " " + str(split_size) not in cv_config.sections():
#         section = 'init'
#     else:
#         section = split_type + " " + str(split_size)
# 
#     if cv_config.has_option(section, 'cv_' + str(targets)):
#         loaded_cv = cv_config.get(section, 'cv_' + str(targets))
#     return loaded_cv


import sys
import numpy as np
from moloi.splits.scaffold_split import scaffold_split
from moloi.splits.cluster_split import cluster_split
from moloi.config_processing import cv_splits_save, cv_splits_load
from sklearn.model_selection import train_test_split


def create_cv(smiles, split_type, n_cv, random_state, y=False):
    if split_type == "scaffold":
        count = n_cv
        n_cv = [([], []) for _ in range(count)]
        for i in range(count):

# Path: moloi/splits/scaffold_split.py
# def scaffold_split(smiles, frac_train=.8, seed=777):
#     """
#     Splits compounds into train/validation/test by scaffold.
# 
#     Warning: if there is one very popular scaffold can produce unbalanced split
# 
#     Params
#     ------
#     smiles: list
#         List of smiles
#     frac_train: float, default: 0.8
#         Float in [0, 1] range indicating size of the training set
#     seed: int
#         Used to shuffle smiles before splitting
# 
#     Notes
#     -----
#     Copied from https://github.com/deepchem/deepchem/blob/master/deepchem/splits/splitters.py
#     """
#     smiles = list(smiles)
# 
#     scaffolds = {}
# 
#     for ind, smi in tqdm(enumerate(smiles), total=len(smiles)):
#         scaffold = generate_scaffold(smi)
#         if scaffold not in scaffolds:
#             scaffolds[scaffold] = [ind]
#         else:
#             scaffolds[scaffold].append(ind)
# 
#     scaffolds_keys = list(scaffolds)
#     rng = np.random.RandomState(seed)
#     rng.shuffle(scaffolds_keys)
# 
#     train_cutoff = frac_train * len(smiles)
#     train_inds, test_inds = [], []
#     for scaffold_key in scaffolds_keys:
#         if len(train_inds) > train_cutoff:
#             test_inds += scaffolds[scaffold_key]
#         else:
#             train_inds += scaffolds[scaffold_key]
#     return train_inds, test_inds
# 
# Path: moloi/splits/cluster_split.py
# def cluster_split(smi, test_cluster_id, n_splits):
#     X_molprint, _ = molprint2d_count_fingerprinter(np.array(smi))
#     X_molprint = X_molprint.astype(np.float32)
#     K_molprint_tanimoto = tanimoto_similarity(X_molprint, X_molprint)
# 
#     clustering = BalancedAgglomerativeClustering(n_clusters=n_splits)
#     logger.info("Clustering")
#     labels = clustering.fit_predict(K_molprint_tanimoto)
# 
#     train_ids = np.where(np.array(labels) != test_cluster_id)[0]
#     test_ids = np.where(np.array(labels) == test_cluster_id)[0]
# 
#     return train_ids, test_ids
# 
# Path: moloi/config_processing.py
# def cv_splits_save(split_type, split_size, n_cv, data_config, targets):
#     sv_config = ConfigParser.ConfigParser()
# 
#     sv_config.read(data_config)
#     try:
#         sv_config[split_type + " " + str(split_size)]['cv_' + str(targets)] = str(n_cv)
#     except:
#         with open(data_config, "a") as ini:
#             ini.write('[' + split_type + " " + str(split_size) + ']')
#         sv_config.read(data_config)
#         sv_config[split_type + " " + str(split_size)]['cv_' + str(targets)] = str(n_cv)
#     with open(data_config, 'w') as configfile:
#         sv_config.write(configfile)
# 
# def cv_splits_load(split_type, split_size, data_config, targets):
#     loaded_cv = False
#     cv_config = ConfigParser.ConfigParser()
#     cv_config.read(data_config)
#     if split_type is False:
#         section = 'init'
#     else:
#         section = split_type + " " + str(split_size)
# 
#     if split_type + " " + str(split_size) not in cv_config.sections():
#         section = 'init'
#     else:
#         section = split_type + " " + str(split_size)
# 
#     if cv_config.has_option(section, 'cv_' + str(targets)):
#         loaded_cv = cv_config.get(section, 'cv_' + str(targets))
#     return loaded_cv


# Path: moloi/splits/cv.py
import sys
import numpy as np
from moloi.splits.scaffold_split import scaffold_split
from moloi.splits.cluster_split import cluster_split
from moloi.config_processing import cv_splits_save, cv_splits_load
from sklearn.model_selection import train_test_split



def create_cv(smiles, split_type, n_cv, random_state, y=False):
    if split_type == "scaffold":
        count = n_cv
        n_cv = [([], []) for _ in range(count)]
        for i in range(count):

            train, test = scaffold_split(smiles, frac_train=1 - ((len(smiles)/count)/(len(smiles)/100))/100,