from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit import RDConfig
import os
from rdkit.Chem.Draw import IPythonConsole
from rdkit.Chem import Draw

cdk2mols = [x for x in Chem.SDMolSupplier('data/cdk2.sdf',removeHs=False) if x is not None]

cdk2mols[0]

cdk2mols[1]

from rdkit.Chem import PyMol
v = PyMol.MolViewer()

v.ShowMol(cdk2mols[0])
v.GetPNG(h=300)

v.ShowMol(cdk2mols[1],name=cdk2mols[1].GetProp('_Name'))   # <- this replaces the current molecule
v.GetPNG(h=300)

v.ShowMol(cdk2mols[2],showOnly=False,name=cdk2mols[2].GetProp('_Name'))   # <- this replaces the current molecule
v.GetPNG(h=300)

import copy
tm = copy.deepcopy(cdk2mols[0])
tm.RemoveAllConformers()
tm.GetNumConformers()

cids=AllChem.EmbedMultipleConfs(tm,pruneRmsThresh=1.0)

for cid in cids: AllChem.UFFOptimizeMolecule(tm,confId=cid)

v.ShowMol(tm,name='conf-0',confId=cids[0])
v.GetPNG(h=300)

v.ShowMol(tm,name='conf-1',confId=cids[1],showOnly=False)
v.GetPNG(h=300)

tm

core = Chem.MolFromSmarts('c12ncnc1ncnc2')
match=tm.GetSubstructMatch(core)
tm

match

AllChem.AlignMolConformers(tm,atomIds=match)

v.ShowMol(tm,name='conf-1',confId=cids[1],showOnly=False)  # <- the same name we already used, so PyMol replaces the object
v.GetPNG(h=300)

v.DeleteAll()
for cid in cids: 
    v.ShowMol(tm,confId=cid,name='Conf-%d'%cid,showOnly=False)
v.GetPNG(h=300)

conf = tm.GetConformer()   # <- gives the default (first) conformer
conf.GetAtomPosition(0)

list(conf.GetAtomPosition(0))

m1 = cdk2mols[0]
m2 = cdk2mols[1]
v.ShowMol(m1,name=m1.GetProp('_Name'))
v.ShowMol(m2,name=m2.GetProp('_Name'),showOnly=False)
v.GetPNG(h=300)

core

match1 = m1.GetSubstructMatch(core)
match2= m2.GetSubstructMatch(core)

AllChem.AlignMol(m2,m1,atomMap=zip(match2,match1))   # <- m2 is aligned to m1, return value is the RMSD for the alignment

v.ShowMol(m1,name=m1.GetProp('_Name'))
v.ShowMol(m2,name=m2.GetProp('_Name'),showOnly=False)
v.GetPNG(h=300)

from rdkit.Chem import MCS

m1 = cdk2mols[0]
m2 = cdk2mols[2]
m1

m2

mcs = MCS.FindMCS([m1,m2],ringMatchesRingOnly=True)
mcs

core = Chem.MolFromSmarts(mcs.smarts)
core

match1 = m1.GetSubstructMatch(core)
match2= m2.GetSubstructMatch(core)
AllChem.AlignMol(m2,m1,atomMap=zip(match2,match1))   # <- m2 is aligned to m1

v.ShowMol(m1,name=m1.GetProp('_Name'))
v.ShowMol(m2,name=m2.GetProp('_Name'),showOnly=False)
v.GetPNG(h=300)

m1

m2

ffact = AllChem.BuildFeatureFactory(os.path.join(RDConfig.RDDataDir,'BaseFeatures.fdef'))

ffact.GetFeatureFamilies()

m = cdk2mols[0]
feats = ffact.GetFeaturesForMol(m)
[x.GetFamily() for x in feats]

d1 = feats[0]

d1.GetAtomIds()

list(d1.GetPos())

list(m.GetConformer().GetAtomPosition(10))

v.ShowMol(m,name=m.GetProp('_Name'))

donors = ffact.GetFeaturesForMol(m,includeOnly='Donor')

for donor in donors: 
    v.server.sphere(list(donor.GetPos()),0.5,(1,0,1),'donors',1)
v.GetPNG(h=300)

defs=ffact.GetFeatureDefs()

sorted(defs.keys())

defs['Donor.SingleAtomDonor']

for feat in ffact.GetFeaturesForMol(m,includeOnly='PosIonizable'): 
    v.server.sphere(list(feat.GetPos()),0.5,(1,0,0),'Pos',1) 
v.GetPNG(h=300)

defs['PosIonizable.Imidazole']

print >>file('data/output.mol','w+'),Chem.MolToMolBlock(m)

!python $RDBASE/rdkit/Chem/Features/ShowFeats.py data/output.mol

v.GetPNG()

cdk2mols[0]

cdk2mols[1]

core = Chem.MolFromSmarts('c12ncn([#1])c1nc(N)nc2OC')
core

em = Chem.EditableMol(cdk2mols[0])
match = cdk2mols[0].GetSubstructMatch(core)
for idx in range(cdk2mols[0].GetNumAtoms()-1,-1,-1): 
    if idx not in match:
        em.RemoveAtom(idx)
coreMol = em.GetMol()
Chem.SanitizeMol(coreMol)
coreMol

newMol = Chem.MolFromSmiles('c12nc[nH]c1nc(N)nc2OCc1cc(C)ccc1')
newMol

newMol=Chem.AddHs(newMol)
newMol3D=AllChem.ConstrainedEmbed(newMol,coreMol)

v.ShowMol(cdk2mols[0],name='reference')
v.ShowMol(newMol3D,name='new mol',showOnly=False)
v.GetPNG(h=300)

from rdkit.Chem.Subshape import SubshapeBuilder,SubshapeAligner,SubshapeObjects

Draw.MolsToGridImage(cdk2mols[:20],molsPerRow=4)

ref = Chem.Mol(cdk2mols[0].ToBinary())
probe = Chem.Mol(cdk2mols[1].ToBinary())
v.ShowMol(ref,name='ref')
v.ShowMol(probe,name='probe',showOnly=False)
v.GetPNG()

AllChem.CanonicalizeConformer(ref.GetConformer())
builder = SubshapeBuilder.SubshapeBuilder()
builder.gridDims=(20.,20.,10)
builder.gridSpacing=0.5
builder.winRad=4.

refShape = builder.GenerateSubshapeShape(ref)
probeShape = builder(probe,terminalPointsOnly=True)

v.ShowMol(ref,name='ref')
SubshapeObjects.DisplaySubshape(v,refShape,'ref_shape')
v.server.do('set transparency=0.5')
v.GetPNG()

aligner = SubshapeAligner.SubshapeAligner()

algs = aligner.GetSubshapeAlignments(ref,refShape,probe,probeShape,builder)
algs

alg = algs[0]
AllChem.TransformMol(probe,alg.transform)
v.ShowMol(ref,name='ref')
v.ShowMol(probe,name='probe',showOnly=False)
v.GetPNG()

probe = Chem.Mol(cdk2mols[8].ToBinary())
v.ShowMol(ref,name='ref')
v.ShowMol(probe,name='probe',showOnly=False)
v.GetPNG()

probeShape = builder(probe,terminalPointsOnly=True)
aligner = SubshapeAligner.SubshapeAligner()

algs = aligner.GetSubshapeAlignments(ref,refShape,probe,probeShape,builder)
algs

alg = algs[0]
AllChem.TransformMol(probe,alg.transform)
v.ShowMol(ref,name='ref')
v.ShowMol(probe,name='probe',showOnly=False)
v.GetPNG()

probe = Chem.Mol(cdk2mols[8].ToBinary())
alg = algs[1]
AllChem.TransformMol(probe,alg.transform)
v.ShowMol(ref,name='ref')
v.ShowMol(probe,name='probe',showOnly=False)
v.GetPNG()

probe = Chem.Mol(cdk2mols[8].ToBinary())
alg = algs[2]
AllChem.TransformMol(probe,alg.transform)
v.ShowMol(ref,name='ref')
v.ShowMol(probe,name='probe',showOnly=False)
v.GetPNG()

probe = Chem.Mol(cdk2mols[8].ToBinary())
algs = SubshapeAligner.ClusterAlignments(probe,algs,builder,neighborTol=0.15)
algs

probe = Chem.MolFromSmiles('c12nc[nH]c1nc(N)nc2OC')
probe = Chem.AddHs(probe)
probe

AllChem.EmbedMolecule(probe)
AllChem.UFFOptimizeMolecule(probe)

probeShape = builder(probe,terminalPointsOnly=True)
aligner = SubshapeAligner.SubshapeAligner()

algs = aligner.GetSubshapeAlignments(ref,refShape,probe,probeShape,builder)
algs = SubshapeAligner.ClusterAlignments(probe,algs,builder,neighborTol=0.15)
algs

for i,alg in enumerate(algs):
    print i,alg.dirMatch,alg.shapeDist,alg.triangleSSD

tprobe = Chem.Mol(probe.ToBinary())
alg = algs[9]
AllChem.TransformMol(tprobe,alg.transform)
v.ShowMol(ref,name='ref')
v.ShowMol(tprobe,name='probe',showOnly=False)
v.GetPNG()

builder2 = SubshapeBuilder.SubshapeBuilder()
builder2.gridDims=(20.,20.,10)
builder2.gridSpacing=0.25
builder2.winRad=4.

refShape2 = builder2.GenerateSubshapeShape(ref)

v.ShowMol(ref,name='ref')
SubshapeObjects.DisplaySubshape(v,refShape2,'ref_shape')
v.server.do('set transparency=0.5')
v.GetPNG()

probeShape = builder(probe,terminalPointsOnly=True)
aligner = SubshapeAligner.SubshapeAligner()

algs = aligner.GetSubshapeAlignments(ref,refShape2,probe,probeShape,builder2)
algs = SubshapeAligner.ClusterAlignments(probe,algs,builder2,neighborTol=0.15)
algs

for i,alg in enumerate(algs):
    print i,alg.dirMatch,alg.shapeDist,alg.triangleSSD

tprobe = Chem.Mol(probe.ToBinary())
alg = algs[2]
AllChem.TransformMol(tprobe,alg.transform)
v.ShowMol(ref,name='ref')
v.ShowMol(tprobe,name='probe',showOnly=False)
v.GetPNG()