#!usr/bin/env python

# 2015-08-27	(jmg)	test to set a given double bond to STEREOANY

from rdkit import Chem, rdBase
from rdkit.Chem import AllChem

molblock="""TEST_in
     RDKit          2D

 13 13  0  0  0  0  0  0  0  0999 V2000
   -0.9060    1.3112    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
   -1.6205    0.8987    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
   -1.6205    0.0737    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
   -0.9060   -0.3389    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
   -0.1916    0.0737    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
   -0.1916    0.8987    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
   -0.9060   -1.1639    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
   -0.1916   -1.5764    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
    0.5229    1.3112    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
    0.5229    2.1362    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
    1.2374    2.5487    0.0000 S   0  0  0  0  0  0  0  0  0  0  0  0
    0.5229   -1.1639    0.0000 O   0  0  0  0  0  0  0  0  0  0  0  0
   -0.1916    2.5487    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
  1  2  2  0
  2  3  1  0
  3  4  2  0
  4  5  1  0
  5  6  2  0
  1  6  1  0
  4  7  1  0
  7  8  2  0
  6  9  1  0
  9 10  2  0
 10 11  1  0
  8 12  1  0
 10 13  1  0
M  END

$$$$
"""

print "RDKIT", rdBase.rdkitVersion, "\n"
print "TEST TO SET ONE GIVEN DOUBLE BOND STEREOCHEMISTRY TO 'STEREOANY'\n"

print "\n############################ INPUT STEREOCHEMISTRY ############################"
mol = Chem.MolFromMolBlock(molblock)
name = mol.GetProp("_Name")
print Chem.MolToMolBlock(mol,includeStereo=True)
# BONDS
print "\n###### BONDS ######\n"
print "Idx","\t", "Begin","\t", "End","\t", "Stereo","\t\t", "Dir"
for b in mol.GetBonds():
	print b.GetIdx(),"\t", b.GetBeginAtomIdx(),"\t", b.GetEndAtomIdx(),"\t", b.GetStereo(),"\t", b.GetBondDir()
print "\n###### ATOMS ######\n"	
print "Idx","\t", "Symb","\t" "ChiralTag"
for a in mol.GetAtoms():
	print a.GetIdx(),"\t", a.GetSymbol(),"\t", a.GetChiralTag()

print "\nNo flag in the bond block to specify the double bond stereochemistry.\n"
################################################################################

print "\n########################## EXPLICIT STEREOCHEMISTRY ############################"
mol1 = Chem.MolFromMolBlock(molblock)
mol1.SetProp("_Name", name.replace("in", "allExplicitStereo"))
AllChem.AssignStereochemistry(mol1, cleanIt=True, force=True, flagPossibleStereoCenters=True)
print Chem.MolToMolBlock(mol1,includeStereo=True)
# BONDS
print "\n###### BONDS ######\n"
print "Idx","\t", "Begin","\t", "End","\t", "Stereo","\t\t", "Dir"
for b in mol1.GetBonds():
	print b.GetIdx(),"\t", b.GetBeginAtomIdx(),"\t", b.GetEndAtomIdx(),"\t", b.GetStereo(),"\t", b.GetBondDir()
print "\n###### ATOMS ######\n"	
print "Idx","\t", "Symb","\t" "ChiralTag"
for a in mol1.GetAtoms():
	print a.GetIdx(),"\t", a.GetSymbol(),"\t", a.GetChiralTag()

print "\nStill no flag in the bond block to specify the double bond stereochemistry.\n"
################################################################################

print "\n########################## REMOVE ALL STEREOCHEMISTRY ##########################"
mol2=Chem.MolFromMolBlock(molblock)
Chem.RemoveStereochemistry(mol2)
mol2.SetProp("_Name", name.replace("in", "allStereoRemoved"))
print Chem.MolToMolBlock(mol2, includeStereo=True)

print "\n###### BONDS ######\n"
print "Idx","\t", "Begin","\t", "End","\t", "Stereo","\t\t", "Dir"
for b in mol2.GetBonds():
	print b.GetIdx(),"\t", b.GetBeginAtomIdx(),"\t", b.GetEndAtomIdx(),"\t", b.GetStereo(),"\t", b.GetBondDir()
print "\n###### ATOMS ######\n"	
print "Idx","\t", "Symb","\t" "ChiralTag"
for a in mol2.GetAtoms():
	print a.GetIdx(),"\t", a.GetSymbol(),"\t", a.GetChiralTag()
	
print "\nThese are the flags I want in the bond block, but the BondStereo is set to STEREONONE instead of STEREOANY.\n"

################################################################################
	
print "\n########################### REMOVE 8-9 DB STEREOCHEMISTRY ##########################"
mol3=Chem.MolFromMolBlock(molblock)
mol3.SetProp("_Name", name.replace("in", "OneStereoDBRemoved?"))
print Chem.MolToMolBlock(mol3, includeStereo=True)

# indices of the atoms making the stereo double bond
idxA0 = 8
idxA1 = 9
b = mol3.GetBondBetweenAtoms(idxA0, idxA1)


#print 
b.SetBondDir(Chem.BondDir.UNKNOWN)
AllChem.AssignStereochemistry(mol3, cleanIt=True, force=True, flagPossibleStereoCenters=True)

print "\n###### BONDS ######\n"
print "Idx","\t", "Begin","\t", "End","\t", "Stereo","\t\t", "Dir"
for b in mol3.GetBonds():
	print b.GetIdx(),"\t", b.GetBeginAtomIdx(),"\t", b.GetEndAtomIdx(),"\t", b.GetStereo(),"\t", b.GetBondDir()
print "\n###### ATOMS ######\n"	
print "Idx","\t", "Symb","\t" "ChiralTag"
for a in mol3.GetAtoms():
	print a.GetIdx(),"\t", a.GetSymbol(),"\t", a.GetChiralTag()

print "\nI don't get any 'cis or trans' flags in the bond block and the other stereo db BondDir is changed.\n"

################################################################################
print ""