#! /usr/bin/env python
# -*- coding: utf-8 -*-
# génération à partir du code smile d'alcools ddes caractéristiques 3D
# et de fichiers .pdb lisible par Jmol
# références : 
# http://openbabel.org/wiki/Python
# http://openbabel.org/docs/current/UseTheLibrary/PythonDoc.html#using-iterators
# http://pythonchem.blogspot.be/2012/08/fun-and-little-disappointment-with.html
# https://gist.github.com/cstein/3294891
# http://nullege.com/codes/search/openbabel

import openbabel

def OBMolMinimize(mol):
    """Minimize a molecule
    """
    ff = openbabel.OBForceField.FindForceField("MMFF94")
    ff.Setup(mol)
    ff.ConjugateGradients(100, 1.0e-5)
    ff.GetCoordinates(mol)
    return mol

def OBStructureFromSmiles(smilesstring, filename=None):
    mol = openbabel.OBMol()
    obConversion = openbabel.OBConversion()
    obConversion.SetInAndOutFormats("smi", "pdb")
    obConversion.ReadString(mol, smilesstring)
    mol.AddHydrogens()
    builder = openbabel.OBBuilder()
    builder.Build(mol)
    mol = OBMolMinimize(mol)
    if filename is None: return mol
    # save structures in subfolder molecules
    obConversion.WriteFile(mol, "molecules/%s.pdb" % filename)

def getAlcohols():
    molecules = dict()
    molecules['methanol'] = 'CO'
    molecules['ethanol'] = 'CCO'
    molecules['1-propanol'] = 'CCCO'
    molecules['isopropanol'] = 'CC(O)C'
    molecules['n-butanol'] = 'CCCCO'
    molecules['butan-2-ol'] = 'CC(O)CC'
    molecules['isobutanol'] = 'CC(C)CO'
    molecules['tert-butanol'] = 'CC(C)(C)O'
    molecules['1-Pentanol'] = 'CCCCCO'
    return molecules

if __name__ == '__main__':

    etab = openbabel.OBElementTable()
    dico = getAlcohols()
    for elem in dico:
        OBStructureFromSmiles(dico[elem], elem)

    # Determine the charge of a molecule
    for file_in in dico:
        obConversion = openbabel.OBConversion()
        obConversion.SetInFormat("pdb")
        mol = openbabel.OBMol()
        obConversion.ReadFile(mol, "molecules/%s.pdb" % file_in)
        # get the charges of the atoms
        charge_model = openbabel.OBChargeModel.FindType("MMFF94")
        charge_model.ComputeCharges(mol)
        partial_charges = charge_model.GetPartialCharges()
        print file_in, '- Formula = ', mol.GetFormula()
        print 'mol wt = ', mol.GetMolWt(), '-  Numb atoms = ', mol.NumAtoms(), '-  Numb bonds = ', mol.NumBonds()
        for atom in openbabel.OBMolAtomIter(mol):
            z=atom.GetAtomicNum()
            v3=atom.GetVector()
            print atom.GetAtomicMass(), z, etab.GetSymbol(z), 'x=',v3.GetX(), 'y=',v3.GetY(), 'z=',v3.GetZ()
        print 'partial charges = ', partial_charges
        print 'total charge = ', "%i" % int(sum(partial_charges))
        print '********'