from __future__ import print_function
from rdkit import RDConfig
import os.path
from rdkit.six.moves import cPickle
from rdkit import Chem
from rdkit.Chem import Descriptors

descrs=['BertzCT','Chi0','HallKierAlpha','Ipc','Kappa3','LabuteASA','SMR_VSA1',
        'SMR_VSA10','SMR_VSA2','SMR_VSA3','SMR_VSA4','SMR_VSA5','SMR_VSA6','SMR_VSA7',
        'SMR_VSA8','SMR_VSA9','SlogP_VSA1','SlogP_VSA10','SlogP_VSA11','SlogP_VSA12',
        'SlogP_VSA2','SlogP_VSA3','SlogP_VSA5','SlogP_VSA6','SlogP_VSA7','SlogP_VSA8',
        'SlogP_VSA9','TPSA','NumHAcceptors','NumHDonors','NumHeteroatoms',
        'NumRotatableBonds','MolLogP']

def runIt(inFileName,outFileName,smiCol=0,maxMols=-1,delim=','):
  outF = open(outFileName,'w+')
  outF.write(','.join(['SMILES']+descrs))
  outF.write('\n')
  mols = []
  nDone = 0
  for line in inD:
    if line[0] != '#':
      splitL = line.strip().split(delim)
      if not splitL: continue
      smi = splitL[smiCol].strip()
      mol = Chem.MolFromSmiles(smi)
      print(smi)
      if mol:
        vals = []
        for descr in descrs:
          fn = getattr(Descriptors,descr)
          try:
            v = fn(mol)
          except Exception:
            v = 666
          vals.append(v)  
        outF.write(','.join([smi]+['%.4f'%x for x in vals]))
        outF.write('\n')
      nDone += 1
      if maxMols>0 and nDone>=maxMols:
        break
  outF.close()    
    
if __name__ == '__main__':
  if 1:
    # 100 PhysProp compounds
    inD = file(os.path.join(RDConfig.RDCodeDir,'Chem','test_data',
                               'PP_descrs_regress.2.csv'),'r').readlines()
    outFileName = os.path.join(RDConfig.RDCodeDir,'Chem','test_data',
                               'PP_descrs_regress.2.csv')
    runIt(inD,outFileName,smiCol=0,delim=',')

    
