# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

from pytorch3dunet.datasets.pdb import StandardPDBDataset, apbsInput
from pathlib import Path
import os
import prody as pr
from openbabel import openbabel
import subprocess
from potsim2 import PotGrid
import numpy as np

# +
src_data_folder = '/home/lorenzo/deep_apbs/srcData/pdbbind_v2019_refined'
name = '4us3'
tmp_data_folder = 'runs/test_3/tmp'
pdb2pqrPath = '/home/lorenzo/pymol/bin/pdb2pqr'

def remove(fname):
    try:
        os.remove(fname)
    except OSError:
        pass
    
def _processPdb():

    src_pdb_file = f'{src_data_folder}/{name}/{name}_protein.pdb'
    tmp_ligand_pdb_file = Path(tmp_data_folder) / name / f'{name}_ligand.pdb'
    os.makedirs(tmp_ligand_pdb_file.parent, exist_ok=True)
    tmp_ligand_pdb_file = str(tmp_ligand_pdb_file)
    remove(tmp_ligand_pdb_file)

    src_mol_file = f"{src_data_folder}/{name}/{name}_ligand.mol2"

    obConversion = openbabel.OBConversion()
    obConversion.SetInAndOutFormats("mol2", "pdb")

    # remove water molecules (found some pdb files with water molecules)
    structure = pr.parsePDB(src_pdb_file)
    protein = structure.select('protein').toAtomGroup()

    # convert ligand to pdb
    ligand = openbabel.OBMol()

    obConversion.ReadFile(ligand, src_mol_file)
    obConversion.WriteFile(ligand, tmp_ligand_pdb_file)

    # select only chains that are close to the ligand (I love ProDy v2)
    ligand = pr.parsePDB(tmp_ligand_pdb_file)
    lresname = ligand.getResnames()[0]
    complx = ligand + protein

    # select ONLY atoms that belong to the protein
    complx = complx.select(f'same chain as exwithin 7 of resname {lresname}')
    complx = complx.select(f'protein and not resname {lresname}')

    return complx, ligand


def _runApbs(out_dir):
    owd = os.getcwd()
    os.chdir(out_dir)

    apbs_in_fname = "apbs-in"
    input = apbsInput(name=name)

    with open(apbs_in_fname, "w") as f:
        f.write(input)

    # generates dx.gz grid file
    proc = subprocess.Popen(
        ["apbs", apbs_in_fname], stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )
    cmd_out = proc.communicate()
    if proc.returncode != 0:
        raise Exception(cmd_out[1].decode())
    os.chdir(owd)

def _genGrids(structure, ligand):

    out_dir = f'{tmp_data_folder}/{name}'
    dst_pdb_file = f'{out_dir}/{name}_protein_trans.pdb'
    remove(dst_pdb_file)
    pqr_output = f"{out_dir}/{name}.pqr"
    remove(pqr_output)
    grid_fname = f"{out_dir}/{name}_grid.dx.gz"
    remove(grid_fname)

    pr.writePDB(dst_pdb_file, structure)
    # pdb2pqr fails to read pdbs with the one line header generated by ProDy...
    with open(dst_pdb_file, 'r') as fin:
        data = fin.read().splitlines(True)
    with open(dst_pdb_file, 'w') as fout:
        fout.writelines(data[1:])

    proc = subprocess.Popen(
        [
            pdb2pqrPath,
            "--with-ph=7.4",
            "--ff=PARSE",
            "--chain",
            dst_pdb_file,
            pqr_output
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    cmd_out = proc.communicate()
    if proc.returncode != 0:
        raise Exception(cmd_out[1].decode())

    print(f'Running apbs on {name}')
    _runApbs(out_dir)

    # read grid
    grid = PotGrid(dst_pdb_file, grid_fname)

    # ligand mask is a boolean NumPy array, can be converted to int: ligand_mask.astype(int)
    ligand_mask = grid.get_ligand_mask(ligand)

    return grid, ligand_mask


# -

structure, ligand = _processPdb()
grid, ligand_mask = _genGrids(structure, ligand)
out_dir = f'{tmp_data_folder}/{name}'
dst_pdb_file = f'{out_dir}/{name}_protein_trans.pdb'
structure2 = pr.parsePDB(dst_pdb_file)
structure2

structure2


# +
def labelGrid(structure, grid):

    retgrid = np.zeros(shape=grid.grid.shape)

    for i,coord in enumerate(structure.getCoords()):
        x,y,z = coord
        binx = int((x - min(grid.edges[0])) / grid.delta[0])
        biny = int((y - min(grid.edges[1])) / grid.delta[1])
        binz = int((z - min(grid.edges[2])) / grid.delta[2])
        
        retgrid[binx,biny,binz] = 1

    return retgrid

retgrid = labelGrid(structure2,grid)
# -

retgrid.sum()

m = np.array([[[i+j for i in range (3)] for j in range(4) ] for k in range(2)] ).transpose()
m.shape

x = np.array([1,2,3])
x2 = x[:,np.newaxis,np.newaxis]
x2.shape

y = m-x2
y.shape



StarcoderdataJupyterScriptsDedupFiltered