#!/usr/bin/env python
"Removes water (HOH) molecules from a molecule"
import sys
try:
    import MolKit
except ImportError, inst:
    print inst
    print "Please install MolKit - http://mgltools.scripps.edu"
    
def usage():
    "Print helpful usage statement to stdout."
    print __file__ +": no input file\n"
    print "\tUsage: removeHOH.py input molecule supported by MolKit\n"


if len(sys.argv) < 2:
       usage()
       sys.exit(1) #help(sys.exit) for more info

mol = MolKit.Read(sys.argv[1]) #FIXME: check if file sys.argv[1] exists
chainsToRemove = [] #this is needed to remove chains with no residues.
for chain in mol.chains:
    residuesHOH = [] #water residues
    for residue in chain.residues:
        if "HOH" in residue.name:
            #Can't do chain.residues.remove(residue). See for instance http://effbot.org/zone/python-list.htm
            residuesHOH.append(residue)
    for HOH in residuesHOH:
        chain.residues.remove(HOH)
    if not chain.residues:
        chainsToRemove.append(chain)

for chain in chainsToRemove:
    mol.chains.remove(chain)
    
from MolKit.pdbWriter import PdbWriter
writer = PdbWriter()
writer.write('out.pdb',mol, records=['ATOM', 'CONECT', 'HETATM'])

#TODO: Add an option for choosing output file name