import numpy as np
import copy

# Отображение структур
import IPython.display
import ipywidgets
from IPython.display import display,display_svg,SVG,Image

# Open Drug Discovery Toolkit
import oddt
import oddt.docking
import oddt.interactions

# Органика
from rdkit.Chem import Draw
from rdkit.Chem.Draw import IPythonConsole

import pmx # Модуль для манипулирования pdb 

pdb=pmx.Model('HYPCU.B99990002.pdb')

for r in pdb.residues[135:]:
    print r # остатки

# создание объектов белок и лиганда
newpdb = pdb.copy()
for r in newpdb.residues[-3:]:
    newpdb.remove_residue(r)
lig = pdb.copy()
del lig.residues[:-3]

# ищем геометрический центр лиганда
x = []
y = []
z = []
for a in lig.atoms:
    coord = a.x 
    x.append(coord[0])
    y.append(coord[1])
    z.append(coord[2])
    
geom_center = [np.mean(x), np.mean(y), np.mean(z)]
geom_center

newpdb.writePDB("myprot.pdb")
lig.writePDB("lig.pdb")

prot = oddt.toolkit.readfile('pdb','myprot.pdb').next()

prot.OBMol.AddPolarHydrogens()
prot.OBMol.AutomaticPartialCharge()


print 'is it the first mol in 1lmp is protein?',prot.protein,':) and MW of this mol is:', prot.molwt 

smiles = ['O=C(NC1C(O)C(O)C(OC1O)CO)O', 'O=C(NC1C(O)C(O)C(OC1O)CO)[NH3+]',
          'O=C(NC1C(O)C(O)C(OC1O)CO)', 'O=C(NC1C(O)C(O)C(OC1O)CO)(c1ccccc1)','O=C(NC1C(O)C(O)C(OC1O)CO)C(=O)O']
mols= []
images =[]

for s in smiles:
    m = oddt.toolkit.readstring('smi', s)
    if not m.OBMol.Has3D(): 
        m.make3D(forcefield='mmff94', steps=150)
        m.removeh()
        m.OBMol.AddPolarHydrogens()

    mols.append(m)
    ###with print m.OBMol.Has3D() was found that:
    ### deep copy needed to keep 3D , write svg make mols flat
    images.append((SVG(copy.deepcopy(m).write('svg'))))
    
display_svg(*images)

#create docking object
dock_obj= oddt.docking.AutodockVina.autodock_vina(
    protein=prot,size=(20,20,20),center=geom_center,
    executable='/usr/bin/vina',autocleanup=True, num_modes=9)

print dock_obj.tmp_dir
print " ".join(dock_obj.params)

# do it
res = dock_obj.dock(mols,prot)

for i,r in enumerate(res):
    print "%4d%10s%8s%8s" %(i, r.formula, r.data['vina_affinity'],  r.data['vina_rmsd_ub'])

import pandas as pd

hbtotal = []
hbstrict = []
phob = []
stack = []
res_sorted = sorted(res, key = lambda x : float(x.data['vina_affinity']))

for i,r in enumerate(res_sorted):
    int1, int2, strict = oddt.interactions.hbonds(prot,r)
    hbtotal.append(len(int1))
    stack.append(len(oddt.interactions.pi_stacking(prot,r)[0]))
    hbstrict.append(strict.sum())
    ph1, ph2 = oddt.interactions.hydrophobic_contacts(prot,r)
    phob.append(len(ph1))
sort_table = pd.DataFrame({'Formula':[r.formula for r in res_sorted],
                   'Affinity, kcal/mol':[r.data['vina_affinity'] for r in res_sorted],
                   'Total number of h-bonds':hbtotal, 'Strict number of h-bonds':hbstrict, 'Stacking': stack,'Hydrophobic':phob, 
                    'RMSD':[r.data['vina_rmsd_ub'] for r in res_sorted]})

sort_table[:]

for i in (9, 0):
    res[i].write(filename='%s.pdb' % str(i), format='pdb')