#!/usr/bin/env python
"Removes atoms with alternate locations"
import sys

try:
    import MolKit
except ImportError, inst:
    print inst
    print "Please install MolKit - http://mgltools.scripps.edu"
    sys.exit(1)
    
def usage():
    "Print helpful usage statement to stdout."
    print __file__ +": no input file\n"
    print "\tUsage: removeAlt.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
for chain in mol.chains:
    for residue in chain.residues:
        altDict = {} #to keep atoms with 
        for atom in residue.atoms:
            if '@' in atom.name:
                list = atom.name.split('@') #'CB@B'=>['CB', 'B']
                if list[-1] in altDict:
                    altDict[list[-1]].append(atom)
                else:
                    altDict[list[-1]] = [atom]         
        if altDict:
            keys = altDict.keys()
            print "\nResidue %s altLoc:"%residue.name,
            print keys,
            print ' Removing ',
            for key in keys[1:]:
                for atom in altDict[key]:
                    print atom.name,
                    residue.remove(atom)

from MolKit.pdbWriter import PdbWriter
writer = PdbWriter()
writer.write('out.pdb',mol)