#!/usr/bin/env python
#############################################################################
### pdb2lig.py - simplistic PDB ligand extractor.
### Uses the simple heuristic that the ligand is likely to have an
### atomcount smaller than a protein and larger than a salt or water.
### OEChem automatically interprets bonding.
#############################################################################
import os,sys
from openeye.oechem import *

MINATOMS=5
MAXATOMS=100

PROG=os.path.basename(sys.argv[0])
if len(sys.argv) not in (2,3):
  OEThrow.Fatal('syntax: %s <PDBIN> [<OUT>]'%PROG)
elif len(sys.argv)==3:
  oms=oemolostream(sys.argv[2])
else:
  oms=oemolostream('.smi')

ims=oemolistream(sys.argv[1])
pdbfile=os.path.basename(sys.argv[1])

if not ims.GetFormat()==OEFormat_PDB:
  OEThrow.Fatal('Input file not PDB.')

partmol=OECreateOEGraphMol()
for mol in ims.GetOEGraphMols():
  count,partlist=OEDetermineComponents(mol)
  pred=PartPred(partlist)
  for i in range(1,count+1):
    pred.SelectPart(i)
    partmol.Clear()
    OESubsetMol(partmol,mol,pred)
    n=partmol.NumAtoms()
    ## sys.stderr.write("natoms = %d\n"%n)
    if n>MAXATOMS:
      continue  ## probably protein
    elif n<MINATOMS:
      continue  ## probably water or salt
    else:
      title=partmol.GetTitle()
      title=pdbfile+' ligand ('+title+')'
      partmol.SetTitle(title)
      OEWriteMolecule(oms,partmol)