import pickle
from rdkit import Chem
from rdkit.Chem import AllChem

if __name__ == '__main__':
    with open('data/geom_drugs/test_Drugs.pkl', 'rb') as fp:
        list_mol = [AllChem.RemoveHs(m) for m in pickle.load(fp)]
    confs = list_mol[0].GetConformers()
    print([a.GetSymbol() for a in list_mol[0].GetAtoms()])
    print(len(confs))
    print(confs[0].GetPositions())
    exit(123)
    d = {}
    for m in list_mol:
        smiles = Chem.MolToSmiles(m)
        d.setdefault(smiles, []).append(m)
    list_list_mol = list(d.values())
    print([len(v) for v in list_list_mol])
