import deepchem as dc

CheckFeaturizer = {
    ('bace_c', 'logreg'): ['ECFP', 1024],
    ('bace_c', 'tf'): ['ECFP', 1024],
    ('bace_c', 'tf_robust'): ['ECFP', 1024],
    ('bace_c', 'rf'): ['ECFP', 1024],
    ('bace_c', 'kernelsvm'): ['ECFP', 1024],
    ('bace_c', 'irv'): ['ECFP', 1024],
    ('bace_c', 'xgb'): ['ECFP', 1024],
    ('bace_c', 'graphconv'): ['GraphConv', 75],
    ('bace_c', 'dag'): ['GraphConv', 75],
    ('bace_c', 'weave'): ['Weave', 75],
    ('bace_c', 'textcnn'): ['Raw', None],
    ('bace_c', 'mpnn'): ['Weave', [75, 14]],
    ('bbbp', 'logreg'): ['ECFP', 1024],
    ('bbbp', 'tf'): ['ECFP', 1024],
    ('bbbp', 'tf_robust'): ['ECFP', 1024],
    ('bbbp', 'rf'): ['ECFP', 1024],
    ('bbbp', 'kernelsvm'): ['ECFP', 1024],
    ('bbbp', 'irv'): ['ECFP', 1024],
    ('bbbp', 'xgb'): ['ECFP', 1024],
    ('bbbp', 'graphconv'): ['GraphConv', 75],
    ('bbbp', 'dag'): ['GraphConv', 75],
    ('bbbp', 'weave'): ['Weave', 75],
    ('bbbp', 'textcnn'): ['Raw', None],
    ('bbbp', 'mpnn'): ['Weave', [75, 14]],
    ('clintox', 'logreg'): ['ECFP', 1024],
    ('clintox', 'tf'): ['ECFP', 1024],
    ('clintox', 'tf_robust'): ['ECFP', 1024],
    ('clintox', 'rf'): ['ECFP', 1024],
    ('clintox', 'kernelsvm'): ['ECFP', 1024],
    ('clintox', 'irv'): ['ECFP', 1024],
    ('clintox', 'xgb'): ['ECFP', 1024],
    ('clintox', 'graphconv'): ['GraphConv', 75],
    ('clintox', 'dag'): ['GraphConv', 75],
    ('clintox', 'weave'): ['Weave', 75],
    ('clintox', 'textcnn'): ['Raw', None],
    ('clintox', 'mpnn'): ['Weave', [75, 14]],
    ('hiv', 'logreg'): ['ECFP', 1024],
    ('hiv', 'tf'): ['ECFP', 1024],
    ('hiv', 'tf_robust'): ['ECFP', 1024],
    ('hiv', 'rf'): ['ECFP', 1024],
    ('hiv', 'kernelsvm'): ['ECFP', 1024],
    ('hiv', 'irv'): ['ECFP', 1024],
    ('hiv', 'xgb'): ['ECFP', 1024],
    ('hiv', 'graphconv'): ['GraphConv', 75],
    ('hiv', 'dag'): ['GraphConv', 75],
    ('hiv', 'weave'): ['Weave', 75],
    ('hiv', 'textcnn'): ['Raw', None],
    ('hiv', 'mpnn'): ['Weave', [75, 14]],
    ('muv', 'logreg'): ['ECFP', 1024],
    ('muv', 'tf'): ['ECFP', 1024],
    ('muv', 'tf_robust'): ['ECFP', 1024],
    ('muv', 'rf'): ['ECFP', 1024],
    ('muv', 'kernelsvm'): ['ECFP', 1024],
    ('muv', 'irv'): ['ECFP', 1024],
    ('muv', 'xgb'): ['ECFP', 1024],
    ('muv', 'graphconv'): ['GraphConv', 75],
    ('muv', 'siamese'): ['GraphConv', 75],
    ('muv', 'attn'): ['GraphConv', 75],
    ('muv', 'res'): ['GraphConv', 75],
    ('muv', 'weave'): ['Weave', 75],
    ('muv', 'textcnn'): ['Raw', None],
    ('muv', 'mpnn'): ['Weave', [75, 14]],
    ('pcba', 'logreg'): ['ECFP', 1024],
    ('pcba', 'tf'): ['ECFP', 1024],
    ('pcba', 'tf_robust'): ['ECFP', 1024],
    ('pcba', 'irv'): ['ECFP', 1024],
    ('pcba', 'xgb'): ['ECFP', 1024],
    ('pcba', 'graphconv'): ['GraphConv', 75],
    ('pcba', 'weave'): ['Weave', 75],
    ('pcba', 'textcnn'): ['Raw', None],
    ('pcba_146', 'logreg'): ['ECFP', 1024],
    ('pcba_146', 'tf'): ['ECFP', 1024],
    ('pcba_146', 'tf_robust'): ['ECFP', 1024],
    ('pcba_146', 'irv'): ['ECFP', 1024],
    ('pcba_146', 'xgb'): ['ECFP', 1024],
    ('pcba_146', 'graphconv'): ['GraphConv', 75],
    ('pcba_146', 'weave'): ['Weave', 75],
    ('pcba_2475', 'logreg'): ['ECFP', 1024],
    ('pcba_2475', 'tf'): ['ECFP', 1024],
    ('pcba_2475', 'tf_robust'): ['ECFP', 1024],
    ('pcba_2475', 'irv'): ['ECFP', 1024],
    ('pcba_2475', 'xgb'): ['ECFP', 1024],
    ('pcba_2475', 'graphconv'): ['GraphConv', 75],
    ('pcba_2475', 'weave'): ['Weave', 75],
    ('sider', 'logreg'): ['ECFP', 1024],
    ('sider', 'tf'): ['ECFP', 1024],
    ('sider', 'tf_robust'): ['ECFP', 1024],
    ('sider', 'rf'): ['ECFP', 1024],
    ('sider', 'kernelsvm'): ['ECFP', 1024],
    ('sider', 'irv'): ['ECFP', 1024],
    ('sider', 'xgb'): ['ECFP', 1024],
    ('sider', 'graphconv'): ['GraphConv', 75],
    ('sider', 'dag'): ['GraphConv', 75],
    ('sider', 'weave'): ['Weave', 75],
    ('sider', 'siamese'): ['GraphConv', 75],
    ('sider', 'attn'): ['GraphConv', 75],
    ('sider', 'res'): ['GraphConv', 75],
    ('sider', 'textcnn'): ['Raw', None],
    ('sider', 'mpnn'): ['Weave', [75, 14]],
    ('tox21', 'logreg'): ['ECFP', 1024],
    ('tox21', 'tf'): ['ECFP', 1024],
    ('tox21', 'tf_robust'): ['ECFP', 1024],
    ('tox21', 'rf'): ['ECFP', 1024],
    ('tox21', 'kernelsvm'): ['ECFP', 1024],
    ('tox21', 'irv'): ['ECFP', 1024],
    ('tox21', 'xgb'): ['ECFP', 1024],
    ('tox21', 'graphconv'): ['GraphConv', 75],
    ('tox21', 'dag'): ['GraphConv', 75],
    ('tox21', 'weave'): ['Weave', 75],
    ('tox21', 'siamese'): ['GraphConv', 75],
    ('tox21', 'attn'): ['GraphConv', 75],
    ('tox21', 'res'): ['GraphConv', 75],
    ('tox21', 'textcnn'): ['Raw', None],
    ('tox21', 'mpnn'): ['Weave', [75, 14]],
    ('toxcast', 'logreg'): ['ECFP', 1024],
    ('toxcast', 'tf'): ['ECFP', 1024],
    ('toxcast', 'tf_robust'): ['ECFP', 1024],
    ('toxcast', 'rf'): ['ECFP', 1024],
    ('toxcast', 'kernelsvm'): ['ECFP', 1024],
    ('toxcast', 'irv'): ['ECFP', 1024],
    ('toxcast', 'xgb'): ['ECFP', 1024],
    ('toxcast', 'graphconv'): ['GraphConv', 75],
    ('toxcast', 'weave'): ['Weave', 75],
    ('toxcast', 'textcnn'): ['Raw', None],
    ('toxcast', 'mpnn'): ['Weave', [75, 14]],
    ('bace_r', 'tf_regression'): ['ECFP', 1024],
    ('bace_r', 'rf_regression'): ['ECFP', 1024],
    ('bace_r', 'krr'): ['ECFP', 1024],
    ('bace_r', 'xgb_regression'): ['ECFP', 1024],
    ('bace_r', 'graphconvreg'): ['GraphConv', 75],
    ('bace_r', 'dag_regression'): ['GraphConv', 75],
    ('bace_r', 'weave_regression'): ['Weave', 75],
    ('bace_r', 'textcnn_regression'): ['Raw', None],
    ('chembl', 'tf_regression'): ['ECFP', 1024],
    ('chembl', 'rf_regression'): ['ECFP', 1024],
    ('chembl', 'krr'): ['ECFP', 1024],
    ('chembl', 'xgb_regression'): ['ECFP', 1024],
    ('chembl', 'graphconvreg'): ['GraphConv', 75],
    ('chembl', 'weave_regression'): ['Weave', 75],
    ('clearance', 'tf_regression'): ['ECFP', 1024],
    ('clearance', 'rf_regression'): ['ECFP', 1024],
    ('clearance', 'krr'): ['ECFP', 1024],
    ('clearance', 'xgb_regression'): ['ECFP', 1024],
    ('clearance', 'graphconvreg'): ['GraphConv', 75],
    ('clearance', 'dag_regression'): ['GraphConv', 75],
    ('clearance', 'weave_regression'): ['Weave', 75],
    ('delaney', 'tf_regression'): ['ECFP', 1024],
    ('delaney', 'rf_regression'): ['ECFP', 1024],
    ('delaney', 'krr'): ['ECFP', 1024],
    ('delaney', 'xgb_regression'): ['ECFP', 1024],
    ('delaney', 'graphconvreg'): ['GraphConv', 75],
    ('delaney', 'dag_regression'): ['GraphConv', 75],
    ('delaney', 'weave_regression'): ['Weave', 75],
    ('delaney', 'mpnn'): ['Weave', [75, 14]],
    ('delaney', 'textcnn_regression'): ['Raw', None],
    ('hopv', 'tf_regression'): ['ECFP', 1024],
    ('hopv', 'rf_regression'): ['ECFP', 1024],
    ('hopv', 'krr'): ['ECFP', 1024],
    ('hopv', 'xgb_regression'): ['ECFP', 1024],
    ('hopv', 'graphconvreg'): ['GraphConv', 75],
    ('hopv', 'dag_regression'): ['GraphConv', 75],
    ('hopv', 'weave_regression'): ['Weave', 75],
    ('lipo', 'tf_regression'): ['ECFP', 1024],
    ('lipo', 'rf_regression'): ['ECFP', 1024],
    ('lipo', 'krr'): ['ECFP', 1024],
    ('lipo', 'xgb_regression'): ['ECFP', 1024],
    ('lipo', 'graphconvreg'): ['GraphConv', 75],
    ('lipo', 'dag_regression'): ['GraphConv', 75],
    ('lipo', 'weave_regression'): ['Weave', 75],
    ('lipo', 'mpnn'): ['Weave', [75, 14]],
    ('lipo', 'textcnn_regression'): ['Raw', None],
    ('nci', 'tf_regression'): ['ECFP', 1024],
    ('nci', 'rf_regression'): ['ECFP', 1024],
    ('nci', 'krr'): ['ECFP', 1024],
    ('nci', 'xgb_regression'): ['ECFP', 1024],
    ('nci', 'graphconvreg'): ['GraphConv', 75],
    ('nci', 'weave_regression'): ['Weave', 75],
    ('ppb', 'tf_regression'): ['ECFP', 1024],
    ('ppb', 'rf_regression'): ['ECFP', 1024],
    ('ppb', 'krr'): ['ECFP', 1024],
    ('ppb', 'xgb_regression'): ['ECFP', 1024],
    ('ppb', 'graphconvreg'): ['GraphConv', 75],
    ('ppb', 'dag_regression'): ['GraphConv', 75],
    ('ppb', 'weave_regression'): ['Weave', 75],
    ('sampl', 'tf_regression'): ['ECFP', 1024],
    ('sampl', 'rf_regression'): ['ECFP', 1024],
    ('sampl', 'krr'): ['ECFP', 1024],
    ('sampl', 'xgb_regression'): ['ECFP', 1024],
    ('sampl', 'graphconvreg'): ['GraphConv', 75],
    ('sampl', 'dag_regression'): ['GraphConv', 75],
    ('sampl', 'weave_regression'): ['Weave', 75],
    ('sampl', 'mpnn'): ['Weave', [75, 14]],
    ('sampl', 'textcnn_regression'): ['Raw', None],
    ('kaggle', 'tf_regression'): [None, 14293],
    ('kaggle', 'rf_regression'): [None, 14293],
    ('kaggle', 'krr'): [None, 14293],
    ('pdbbind', 'tf_regression'): ['grid', 2052],
    ('pdbbind', 'rf_regression'): ['grid', 2052],
    ('pdbbind', 'krr'): ['grid', 2052],
    ('pdbbind', 'graphconvreg'): ['GraphConv', 75],
    ('qm7', 'tf_regression'): ['ECFP', 1024],
    ('qm7', 'rf_regression'): ['ECFP', 1024],
    ('qm7', 'krr'): ['ECFP', 1024],
    ('qm7', 'krr_ft'): [dc.feat.CoulombMatrix(23), 1024],
    ('qm7', 'textcnn_regression'): ['Raw', None],
    ('qm7', 'graphconvreg'): ['GraphConv', 75],
    ('qm7', 'weave_regression'): ['Weave', 75],
    ('qm7', 'tf_regression_ft'): [dc.feat.CoulombMatrix(23), [23, 23]],
    ('qm7', 'dtnn'): [dc.feat.CoulombMatrix(23), [23, 23]],
    ('qm7', 'ani'): ['BPSymmetryFunctionInput', [23, 4]],
    ('qm8', 'tf_regression'): ['ECFP', 1024],
    ('qm8', 'rf_regression'): ['ECFP', 1024],
    ('qm8', 'krr'): ['ECFP', 1024],
    ('qm8', 'graphconvreg'): ['GraphConv', 75],
    ('qm8', 'tf_regression_ft'): [dc.feat.CoulombMatrix(26), [26, 26]],
    ('qm8', 'krr_ft'): [dc.feat.CoulombMatrix(26), 1024],
    ('qm8', 'dtnn'): [dc.feat.CoulombMatrix(26), [26, 26]],
    ('qm8', 'ani'): ['BPSymmetryFunctionInput', [26, 4]],
    ('qm8', 'mpnn'): ['MP', [70, 8]],
    ('qm8', 'weave_regression'): ['Weave', 75],
    ('qm8', 'textcnn_regression'): ['Raw', None],
    ('qm9', 'tf_regression'): ['ECFP', 1024],
    ('qm9', 'rf_regression'): ['ECFP', 1024],
    ('qm9', 'krr'): ['ECFP', 1024],
    ('qm9', 'graphconvreg'): ['GraphConv', 75],
    ('qm9', 'tf_regression_ft'): [dc.feat.CoulombMatrix(29), [29, 29]],
    ('qm9', 'krr_ft'): [dc.feat.CoulombMatrix(29), 1024],
    ('qm9', 'dtnn'): [dc.feat.CoulombMatrix(29), [29, 29]],
    ('qm9', 'ani'): ['BPSymmetryFunctionInput', [29, 4]],
    ('qm9', 'mpnn'): ['MP', [70, 8]],
    ('qm9', 'weave_regression'): ['Weave', 75],
    ('qm9', 'textcnn_regression'): ['Raw', None]
}

CheckSplit = {
    'bace_c': ['random', 'scaffold'],
    'bace_r': ['random', 'scaffold'],
    'bbbp': ['random', 'scaffold'],
    'chembl': ['index', 'random', 'scaffold', 'year'],
    'clearance': ['index', 'random', 'scaffold'],
    'clintox': ['index', 'random', 'scaffold'],
    'delaney': ['index', 'random', 'scaffold'],
    'hiv': ['index', 'random', 'scaffold', 'butina'],
    'hopv': ['index', 'random', 'scaffold', 'butina'],
    'kaggle': ['index'],  # already splitted, no splitter required
    'lipo': ['index', 'random', 'scaffold'],
    'muv': ['index', 'random', 'scaffold', 'task'],
    'nci': ['index', 'random', 'scaffold'],
    'pcba': ['index', 'random', 'scaffold'],
    'pcba_146': ['index', 'random', 'scaffold'],
    'pcba_2475': ['index', 'random', 'scaffold'],
    'pdbbind': ['index', 'random', 'time'],
    'ppb': ['index', 'random', 'scaffold'],
    'qm7': ['index', 'random', 'stratified'],
    'qm8': ['index', 'random', 'stratified'],
    'qm9': ['index', 'random', 'stratified'],
    'sampl': ['index', 'random', 'scaffold'],
    'sider': ['index', 'random', 'scaffold', 'task'],
    'tox21': ['index', 'random', 'scaffold', 'butina', 'task'],
    'toxcast': ['index', 'random', 'scaffold']
}
