from data.qm9.load_qm9 import load_qm9
from data.load_data import SupportedDatasets
from visualize.find_good_mol import find_good_mol

MAX_NUM = -1

if __name__ == '__main__':
    mols, _ = load_qm9(max_num=MAX_NUM)
    find_good_mol(
        dataset_name=SupportedDatasets.QM9,
        list_mol=mols,
        list_special_config=[
            {
                'GENERATE_TYPE': 'rdkit',
                'DERIVE_TYPE': 'newton',
                'COMPARE_TYPE': 'equiv-trunc',
            }, {
                'GENERATE_TYPE': 'rdkit',
                'DERIVE_TYPE': 'newton',
                'COMPARE_TYPE': 'adj3',
            }, {
                'GENERATE_TYPE': 'rdkit',
                'DERIVE_TYPE': 'newton',
                'COMPARE_TYPE': 'kabsch',
            },
        ],
        use_cuda=False
    )
