#!/usr/local/bin/py

# pdbsurf.py
#
# Combines data from exlogp.perl and a HIN file to 
# produce a PDB file with hydrophobicity values in
# the B factor column

# Rajarshi Guha December 2002
# <rajarshi@presidency.com> <http://jijo.cjb.net>
#
#
import string, math, os, getopt, sys
import tempfile


def usage():
    print """
    Usage: pdbsurf.py [OPTIONS]

    -h,--help        This message
    -d,--data        The output file from exlogp.perl
    -m,--molecule    The hin file
    """
    sys.exit(0)

def load_data(filename = None):

    # hdata[] is a list of tuples of the form
    # (atom num,atom type, value)
    # atom num corresponds to the values in the hin file
    hdata = []
    
    if not filename:
        print """
        Error loading the data file 
        """
        sys.exit(0)
    
    try:
        f = open(filename,'r')
    except IOError:
        print """
        Error loading the data file
        """
        sys.exit(0)
    
    f.readline()
    f.readline()

    for line in f:
        if line[0] == '\n': continue

        line = string.split(line)
        hdata.append( (int(line[0]), int(line[1]), float(line[2])) )
    
    f.close()
    return hdata

def load_hin(filename = None):
    # mdata[] is a list of tuples of the form
    # (atom num, atom sym, x,y,z)
    mdata = []
    if not filename:
        print """
        Could'nt load the hin file
        """
        sys.exit(0)
    try:
        f = open(filename,'r')
    except IOError:
        print """
        Could'nt load the hin file
        """
        sys.exit(0)

    for line in f:
        if string.find(line, 'atom') >= 0:
            line = string.split(line)
            mdata.append( (int(line[1]), line[3], float(line[7]), float(line[8]), float(line[9])) )
    f.close()
    return mdata
    
        
        
        
            
if __name__ == '__main__':

    datafile = ''
    hinfile = ''

    if len(sys.argv) < 3:
        usage()

    try:
        opt, args = getopt.getopt(sys.argv[1:], "d:m:h", ["data=","molecule=","help"])
    except getopt.GetoptError:
        usage()
    for o,a in opt:
        if o in ('-h','--help'): 
            usage()
        if o in ('-d', '--data'):
            datafile = a
        if o in ('-m','--molecule'):
            hinfile = a

    ############3

    hdata = load_data(datafile)
    mdata = load_hin(hinfile)


    # Combine the hydropbicity data and the mnolecule info
    # It will be in a list, each element contains the info
    # for a given point (ie atom). The format of each element is
    #
    # (anum, asym, atype, x,y,x, aval)
    #
    combo = []
    for i in range(0, len(hdata)):
        anum, asym, x,y,z = mdata[i]
        anum, atype, aval = hdata[i]
        combo.append( (anum, asym, atype, x,y,z,aval) )
    
    # Convert the hin to a pdb
    pdbfile = string.split(hinfile,'.')[0]+'.pdb'
    os.system('babel -ihin '+hinfile+' -opdb '+pdbfile)
    tfile = tempfile.mktemp()
    # Prepare to write the surface data        
    fin = open(pdbfile,'r')
    fout = open(tfile,'w')

    c = 0
    for line in fin:
        if string.find(line,'ATOM') == 0:
            anum,asym,atype,x,y,z,aval = combo[c]
            line = line[:62] + str(round(aval,2)) + '\n'

            c += 1
            fout.write(line)
        else:
            fout.write(line)
    fin.close()            
    fout.close()

    os.system('mv '+tfile+' '+pdbfile)
    print 'Converted the HIN file to a PDB. Now view with PyMOL'