#!/usr/bin/python
# fixhisname.py
usage='''
usage: $0 in.pdb out.pdb
  This script is used to fix his name :
  To modify HIS to : HID  if only exist HD1
                     HIE  if only exist HE2
                     HIP  if both HD1 and HE2 exist
  by Qiuyu Fu 2014
'''
import sys,os

# Read in parameters
if len(sys.argv)<3:
    print usage
    sys.exit()
infile,outfile=sys.argv[1:3]

# Find HIS first (include HID ,HIE,HIP . for check use)
hisresi=dict()
for line in os.popen("""grep ATOM %s |grep -v REMARK|grep -E 'HIS|HID|HIE|HIP' |cut -b 22-26|uniq """%(infile,)):
    hisresi[line[:-1]]= set()
print hisresi

# Determine HIS type (scan infile first time)
ifp=open(infile,"r")
for line in ifp:
    lineitems=line.split()
    if lineitems[0]!="ATOM":continue
#    chain=line[21].strip()
#    index=line[22:26].strip()
#    resiindex="%s %s"%(chain,index)
    resiindex=line[21:26]
    if not hisresi.has_key(resiindex):continue
    atomtype=line[12:16].strip()
    if atomtype=="HD1":
        hisresi[resiindex].add("HD1")
    if atomtype=="HE2":
        hisresi[resiindex].add("HE2")
ifp.close()
print hisresi

# trans "HD1HE" to "HIP"
for key in hisresi:
    if hisresi[key]==set(("HD1","HE2")):
        hisresi[key]="HIP"
    elif hisresi[key]==set(("HD1",)):
        hisresi[key]="HID"
    elif hisresi[key]==set(("HE2",)):
        hisresi[key]="HIE"
    else:
        hisresi[key]="HIS"
        
    

# Reload in.pdb and correct HIS name and write in to out.pdb
ifp=open(infile,"r")
ofp=open(outfile,"w")
for line in ifp:
    lineitems=line.split()
    if lineitems[0]!="ATOM":
        ofp.write(line)
    else:
#        chain=line[21].strip()
#        index=line[22:26].strip()
#        resiindex="%s %s"%(chain,index)
        resiindex=line[21:26]
        if not hisresi.has_key(resiindex):
            ofp.write(line)
        else:
            histype=hisresi[resiindex]
#            print histype
            ofp.write(line.replace("HIS",histype))
ifp.close()
ofp.close()
    
