from rdkit import Chem
from rdkit import rdBase
from rdkit import RDConfig
import os

from rdkit.Chem import rdMolDescriptors as rdMD
from rdkit.Chem import AllChem

import time

def get3D(m,is3d):
    if not is3d:
        m = Chem.AddHs(m)
        # define the new code from RDKit Molecule 3D ETKDG.
        ps = AllChem.ETKDG()
        ps.randomSeed = 0xf00d
        AllChem.EmbedMolecule(m,ps)
    r= rdMD.CalcAUTOCORR3D(m)+rdMD.CalcRDF(m)+rdMD.CalcMORSE(m)+rdMD.CalcWHIM(m)+rdMD.CalcGETAWAY(m, precision=0.001)
    return r

def generateAll():
    filename='/Users/GVALMTGG/Github/rdkit_mine/Code/GraphMol/Descriptors/test_data/PBF_egfr.sdf'
    suppl = Chem.SDMolSupplier(filename,removeHs=False)
    mols = [x for x in suppl]
    start = time.time()
    for m in mols:
        r= get3D(m,True)
        print(r)
    end = time.time()
    print end-start

def simple_case():
    start = time.time()
    smi = 'CCC(C)COCCCC'
    m = Chem.MolFromSmiles(smi)
    T = get3D(m,False)
    print T
    end = time.time()
    print end-start

if(__name__=='__main__'):
    # FIX: We need to actually add some tests here, but this doees not need to
    # to be done until the C++ code and tests are straightened out.
    generateAll();

    start = time.time()
    smi = 'CCC(C)COCCCC'
    m = Chem.MolFromSmiles(smi)
    T = get3D(m,False)
    print T
    end = time.time()
    print end-start