text: #!/usr/bin/env python
  # coding: utf-8
  
  
  import os
  
  import scipy.spatial
  
  import atom3d.util.file as fi
  import atom3d.util.formats as ft
  import get_labels as lab
  
  from rdkit import Chem
  import Bio.PDB
  from Bio.PDB.PDBIO import Select
  from tqdm import tqdm
  import argparse
  
  
  
  def get_ligand(ligfile):
      """
      Read ligand from PDB dataset into RDKit Mol.
      :param ligfile: file containing ligand information in SDF format
      :type ligfile: str
      """
      lig=Chem.SDMolSupplier(ligfile)[0]
      # Many SDF in PDBBind do not parse correctly. If SDF fails, try loading the mol2 file instead
      if lig is None:
          print('trying mol2...')
          lig=Chem.MolFromMol2File(ligfile[:-4] + '.mol2')
      if lig is None:
          print('failed')
          return None
      lig = Chem.RemoveHs(lig)
      return lig
  
  
  def get_pocket_res(protein, ligand, dist):
      """
      Given a co-crystallized protein and ligand, extract residues within specified distance of ligand.
  
      :param protein: Biopython object containing receptor protein
      :type protein: Bio.PDB.Structure
      :param ligand: RDKit molecule object containing co-crystallized ligand
      :type ligand: rdkit.Chem.rdchem.Mol
      :param dist: distance cutoff for defining binding pocket, in Angstrom
      :type dist: float
  
      :return key_residues: key binding site residues
      :rtype key_residues:  set containing Bio.PDB.Residue objects
      """
      # get protein coordinates
      prot_atoms = [a for a in protein.get_atoms()]
      prot_coords = [atom.get_coord() for atom in prot_atoms]
  
      # get ligand coordinates
      lig_coords = []
      for i in range(0, ligand.GetNumAtoms()):
          pos = ligand.GetConformer().GetAtomPosition(i)
          lig_coords.append([pos.x, pos.y, pos.z])
  
      kd_tree = scipy.spatial.KDTree(prot_coords)
      key_pts = kd_tree.query_ball_point(lig_coords, r=dist, p=2.0)
      key_pts = set([k for l in key_pts for k in l])
  
      key_residues = set()
      for i in key_pts:
          atom = prot_atoms[i]
          res = atom.get_parent()
          if res.get_resname() == 'HOH':
              continue
          key_residues.add(res)
      return key_residues
  
  
  class PocketSelect(Select):
      """
      Selection class for subsetting protein to key binding residues. This is a subclass of :class:`Bio.PDB.PDBIO.Select`.
      """
      def __init__(self, reslist):
          self.reslist = reslist
      def accept_residue(self, residue):
          if residue in self.reslist:
              return True
          else:
              return False
  
  
  def process_files(input_dir):
      """
      Process all protein (pdb) and ligand (sdf) files in input directory.
      
      :param input dir: directory containing PDBBind data
      :type input_dir: str
      
      :return structure_dict: dictionary containing each structure, keyed by PDB code. Each PDB is a dict containing protein as Biopython object and ligand as RDKit Mol object
      :rtype structure_dict: dict
      """
      structure_dict = {}
      pdb_files = fi.find_files(input_dir, 'pdb')
  
      for f in tqdm(pdb_files, desc='pdb files'):
          pdb_id = fi.get_pdb_code(f)
          if pdb_id not in structure_dict:
              structure_dict[pdb_id] = {}
          if '_protein' in f:
              prot = ft.read_any(f)
              structure_dict[pdb_id]['protein'] = prot
  
      lig_files = fi.find_files(input_dir, 'sdf')
      for f in tqdm(lig_files, desc='ligand files'):
          pdb_id = fi.get_pdb_code(f)
          structure_dict[pdb_id]['ligand'] = get_ligand(f)
  
      return structure_dict
  
  
  
  def write_files(pdbid, protein, ligand, pocket, out_path):
      """
      Writes cleaned structure files for protein, ligand, and pocket.
  
      :param pdbid: PDB ID for protein
      :type pdbid: str
      :param protein: object containing proteins
      :type ligand: object containing ligand molecules
      :param pocket: residues contained in pocket, as output by get_pocket_res()
      :type pocket: set containing Bio.PDB.Residue objects
      :param out_path: directory to write files to
      :type out_path: str
      """
      # write protein to mmCIF file
      io = Bio.PDB.MMCIFIO()
      io.set_structure(protein)
      io.save(os.path.join(out_path, f"{pdbid}_protein.cif"))
  
      # write pocket to mmCIF file
      io.save(os.path.join(out_path, f"{pdbid}_pocket.cif"), PocketSelect(pocket))
  
      # write ligand to file
      writer = Chem.SDWriter(os.path.join(out_path, f"{pdbid}_ligand.sdf"))
      writer.write(ligand)
  
  
  def produce_cleaned_dataset(structure_dict, out_path, dist=6.0):
      """
      Generate and save cleaned dataset, given dictionary of structures processed by process_files.
  
      :param structure_dict: dictionary containing protein and ligand structures processed from input files. Should be a nested dict with PDB ID in first level and keys 'protein', 'ligand' in second level.
      :type structure_dict: dict
      :param out_path: path to output directory
      :type out_path: str
      :param dist: distance cutoff for defining binding pocket (in Angstrom), default is 6.0
      :type dist: float, optional
      """
      proteins = []
      pockets = []
      for pdb, data in tqdm(structure_dict.items(), desc='writing to files'):
          protein = structure_dict[pdb]['protein']
          ligand = structure_dict[pdb]['ligand']
          # check for failed ligand (due to bad structure file)
          if ligand is None:
              continue
          pocket_res = get_pocket_res(protein, ligand, dist)
          write_files(pdb, protein, ligand, pocket_res, out_path)
  
  def generate_labels(data_dir, out_dir):
      lab.main(data_dir, out_dir)
  
  
  def main():
      parser = argparse.ArgumentParser()
      parser.add_argument('data_dir', type=str, help='directory where PDBBind is located')
      parser.add_argument('--dist', type=float, default=6.0, help='distance cutoff for defining pocket')
      parser.add_argument('--out_dir', type=str, default=os.getcwd(), help='directory to place cleaned dataset')
      args = parser.parse_args()
  
      if not os.path.exists(args.data_dir):
          raise Exception('Path not found. Please enter valid path to PDBBind dataset.')
  
      if not os.path.exists(args.out_dir):
          os.mkdir(args.out_dir)
  
      structures = process_files(args.data_dir)
      produce_cleaned_dataset(structures, args.out_dir, args.dist)
      generate_labels(args.data_dir, args.out_dir)
  
  
  if __name__ == "__main__":
      main()
  