import os
os.chdir('/public/gkxiao/work/vina/bond_order')

import rdkit
from rdkit import Chem
from rdkit.Chem import AllChem,pdbqt

import py3Dmol

def MolTo3DView(mol, size=(300, 300), style="stick", surface=False, opacity=0.5):
    """Draw molecule in 3D
    
    Args:
    ----
        mol: rdMol, molecule to show
        size: tuple(int, int), canvas size
        style: str, type of drawing molecule
               style can be 'line', 'stick', 'sphere', 'carton'
        surface, bool, display SAS
        opacity, float, opacity of surface, range 0.0-1.0
    Return:
    ----
        viewer: py3Dmol.view, a class for constructing embedded 3Dmol.js views in ipython notebooks.
    """
    assert style in ('line', 'stick', 'sphere', 'carton')
    mblock = Chem.MolToMolBlock(mol)
    viewer = py3Dmol.view(width=size[0], height=size[1])
    viewer.addModel(mblock, 'mol')
    viewer.setStyle({style:{}})
    if surface:
        viewer.addSurface(py3Dmol.SAS, {'opacity': opacity})
    viewer.zoomTo()
    return viewer

qmol = """\
REMARK  Name = 127
REMARK  4 active torsions:
REMARK  status: ('A' for Active; 'I' for Inactive)
REMARK    1  A    between atoms: O_1  and  C_10
REMARK    2  A    between atoms: O_2  and  C_11
REMARK    3  A    between atoms: C_4  and  C_5
REMARK    4  A    between atoms: C_5  and  C_11
REMARK                            x       y       z     vdW  Elec       q    Type
REMARK                         _______ _______ _______ _____ _____    ______ ____
ROOT
ATOM      4  C   UNL     1       0.238  -0.019  -0.559  0.00  0.00    +0.000 A
ATOM      6  C   UNL     1      -0.443  -1.225  -0.335  0.00  0.00    +0.000 A
ATOM      7  C   UNL     1      -0.455   1.192  -0.414  0.00  0.00    +0.000 A
ATOM      8  C   UNL     1      -1.795  -1.222   0.016  0.00  0.00    +0.000 A
ATOM      9  C   UNL     1      -1.806   1.194  -0.064  0.00  0.00    +0.000 A
ATOM     10  C   UNL     1      -2.465  -0.011   0.147  0.00  0.00    +0.000 A
ATOM     14  H   UNL     1       0.077  -2.176  -0.428  0.00  0.00    +0.000 HD
ATOM     15  H   UNL     1       0.052   2.142  -0.568  0.00  0.00    +0.000 HD
ATOM     16  H   UNL     1      -2.301  -2.167   0.186  0.00  0.00    +0.000 HD
ATOM     17  H   UNL     1      -2.341   2.133   0.049  0.00  0.00    +0.000 HD
ENDROOT
BRANCH   4   5
ATOM      5  C   UNL     1       1.702  -0.026  -0.926  0.00  0.00    +0.000 C
ATOM     12  H   UNL     1       1.920   0.839  -1.565  0.00  0.00    +0.000 HD
ATOM     13  H   UNL     1       1.921  -0.929  -1.508  0.00  0.00    +0.000 HD
BRANCH   5  11
ATOM      3  O   UNL     1       2.359   0.085   1.444  0.00  0.00    +0.000 OA
ATOM     11  C   UNL     1       2.635   0.013   0.260  0.00  0.00    +0.000 C
BRANCH  11   2
ATOM      2  O   UNL     1       3.929  -0.043  -0.128  0.00  0.00    +0.000 OA
ATOM     19  H   UNL     1       4.428  -0.009   0.716  0.00  0.00    +0.000 HD
ENDBRANCH  11   2
ENDBRANCH   5  11
ENDBRANCH   4   5
BRANCH  10   1
ATOM      1  O   UNL     1      -3.784   0.035   0.488  0.00  0.00    +0.000 OA
ATOM     18  H   UNL     1      -4.099  -0.874   0.617  0.00  0.00    +0.000 HD
ENDBRANCH  10   1
TORSDOF 4
"""

读入起始的sdf分子，去氢后作为模板

suppl = Chem.SDMolSupplier('EnergyMin_127.sdf', removeHs=False)
smol = suppl[0]

viewer = MolTo3DView(smol, size=(600, 300), style='stick')
viewer.show()
#看看起始结构的样子

You appear to be running in JupyterLab (or JavaScript failed to load for some other reason). You need to install the 3dmol extension:
jupyter labextension install jupyterlab_3dmol

template = Chem.RemoveHs(smol)

viewer = MolTo3DView(template, size=(600, 300), style='stick')
viewer.show()
#看看去氢后的样子

You appear to be running in JupyterLab (or JavaScript failed to load for some other reason). You need to install the 3dmol extension:
jupyter labextension install jupyterlab_3dmol

#看看PDBQT文件
print(qmol)

REMARK  Name = 127
REMARK  4 active torsions:
REMARK  status: ('A' for Active; 'I' for Inactive)
REMARK    1  A    between atoms: O_1  and  C_10
REMARK    2  A    between atoms: O_2  and  C_11
REMARK    3  A    between atoms: C_4  and  C_5
REMARK    4  A    between atoms: C_5  and  C_11
REMARK                            x       y       z     vdW  Elec       q    Type
REMARK                         _______ _______ _______ _____ _____    ______ ____
ROOT
ATOM      4  C   UNL     1       0.238  -0.019  -0.559  0.00  0.00    +0.000 A
ATOM      6  C   UNL     1      -0.443  -1.225  -0.335  0.00  0.00    +0.000 A
ATOM      7  C   UNL     1      -0.455   1.192  -0.414  0.00  0.00    +0.000 A
ATOM      8  C   UNL     1      -1.795  -1.222   0.016  0.00  0.00    +0.000 A
ATOM      9  C   UNL     1      -1.806   1.194  -0.064  0.00  0.00    +0.000 A
ATOM     10  C   UNL     1      -2.465  -0.011   0.147  0.00  0.00    +0.000 A
ATOM     14  H   UNL     1       0.077  -2.176  -0.428  0.00  0.00    +0.000 HD
ATOM     15  H   UNL     1       0.052   2.142  -0.568  0.00  0.00    +0.000 HD
ATOM     16  H   UNL     1      -2.301  -2.167   0.186  0.00  0.00    +0.000 HD
ATOM     17  H   UNL     1      -2.341   2.133   0.049  0.00  0.00    +0.000 HD
ENDROOT
BRANCH   4   5
ATOM      5  C   UNL     1       1.702  -0.026  -0.926  0.00  0.00    +0.000 C
ATOM     12  H   UNL     1       1.920   0.839  -1.565  0.00  0.00    +0.000 HD
ATOM     13  H   UNL     1       1.921  -0.929  -1.508  0.00  0.00    +0.000 HD
BRANCH   5  11
ATOM      3  O   UNL     1       2.359   0.085   1.444  0.00  0.00    +0.000 OA
ATOM     11  C   UNL     1       2.635   0.013   0.260  0.00  0.00    +0.000 C
BRANCH  11   2
ATOM      2  O   UNL     1       3.929  -0.043  -0.128  0.00  0.00    +0.000 OA
ATOM     19  H   UNL     1       4.428  -0.009   0.716  0.00  0.00    +0.000 HD
ENDBRANCH  11   2
ENDBRANCH   5  11
ENDBRANCH   4   5
BRANCH  10   1
ATOM      1  O   UNL     1      -3.784   0.035   0.488  0.00  0.00    +0.000 OA
ATOM     18  H   UNL     1      -4.099  -0.874   0.617  0.00  0.00    +0.000 HD
ENDBRANCH  10   1
TORSDOF 4

将PDBQT转为RDKIT Mol

mol = pdbqt.MolFromPDBQTBlock(qmol, sanitize=False, removeHs=True)
#去氢
mol = Chem.RemoveHs(mol)

#看看3D结构，可以发现键级不对
viewer = MolTo3DView(mol, size=(600, 300), style='stick')
viewer.show()

You appear to be running in JupyterLab (or JavaScript failed to load for some other reason). You need to install the 3dmol extension:
jupyter labextension install jupyterlab_3dmol

#用模板归属键级
mol_corr = AllChem.AssignBondOrdersFromTemplate(template, mol)

#再看3D结构，发现键级对了
viewer = MolTo3DView(mol_corr, size=(600, 300), style='stick')
viewer.show()

You appear to be running in JupyterLab (or JavaScript failed to load for some other reason). You need to install the 3dmol extension:
jupyter labextension install jupyterlab_3dmol