from __future__ import print_function
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem.Draw import IPythonConsole
IPythonConsole.ipython_useSVG=True
from rdkit.Chem.rdRGroupDecomposition import RGroupDecomposition
import pandas as pd
from rdkit.Chem import PandasTools
from collections import OrderedDict
from IPython.display import HTML
from rdkit import rdBase
rdBase.DisableLog("rdApp.debug")

cores = OrderedDict()

cores['six'] = Chem.MolFromSmarts('[#8]-[#6](=O)-*~1~*~*~*~*~2-[#6]-[#6](=O)-*~1~2')
cores['five'] = Chem.MolFromSmarts('[#8]-[#6](=O)-*~1~*~*~*~2-[#6]-[#6](=O)-*~1~2')


from rdkit.Chem import Draw
Draw.MolsToGridImage(cores.values())

"""
Explanation: Example complicated lactam analysis.
The input cores are generic in atoms and bonds, the system figures out the correct cores for grouping purposes.
This may be overkill ( and certainly takes a long enough time ).  What the algorithm is doing is first mapping each core to each other and then trying to optimizes the side chain selection across all sets.  I.e. it tries to make the matching choices such that all R1's are the same, all R2's are the same and so on.
n.b. It is important that the cores are input in most to least specific as the RGroupDecomposition takes the first matching core.
End of explanation
"""


rg = RGroupDecomposition(cores.values())

mols = []
count = 0
maxidx = -1
for line in open("compounds.txt"):
    sm = line.split()[-1]
    m = Chem.MolFromSmiles(sm)
    if m:
        count += 1
        maxidx = max(maxidx,rg.Add(m))
        if count % 1000 == 0:
            print("Added", count)

print ("Added %s to RGroup Decomposition out of %s"%(maxidx, count))

rg.Process()

"""
Explanation: To use RGroupDecomposition:

construct the class on the core rg = RGroupDecomposition(core)
Call rg.Add( mol ) on the molecules.  If this returns -1, the molecule is not
     compatible with the core
After all molecules are added, call rg.Process() to complete the rgroup
     decomposition.
End of explanation
"""


from rdkit import rdBase
rdBase.DisableLog("rdApp.*")

"""
Explanation: Rendering this many molecules is pretty taxing for the current system, this takes a while.
It is useful to disable logging here.  When making RGroup renderings there 
are a lot of sanitization warnings.
End of explanation
"""


frame = pd.DataFrame(rg.GetRGroupsAsColumns())
PandasTools.ChangeMoleculeRendering(frame)

"""
Explanation: The RGroupDecomposition code is quite compatible with the python pandas integration.
Calling rg.GetRGroupsAsColumns() can be sent directly into a pandas table.
n.b. You need to call PandasTools.ChangeMoleculeRendering(frame) to allow the molecules
to be rendered properly.
End of explanation
"""


f2 = pd.DataFrame(frame.head())
PandasTools.ChangeMoleculeRendering(f2)
HTML(f2.to_html())

"""
Explanation: Just show the first few (for speed and to keep the notebook small)
End of explanation
"""


