from os import system
import types

class DSSP:
    def __init__(self,fstr_pdb,dssp="~/bin/dssp"):
	system(dssp+" "+fstr_pdb+" dssp.out 2>temp.junk")
        fp = open("dssp.out","r")
        while(1):
            if(fp.readline()[2] == '#'): break
        self.content = fp.readlines()
	fp.close()

        self.secArray = []
	self.phiArray = []
	self.psiArray = []
	
	import Numeric as N
	rn = len(self.content)+2;
	self.hbond  = N.zeros((rn,rn),N.Float) # NH--O
	self.hbond2 = N.zeros((rn,rn),N.Float) # O--HN
        for line in self.content:
            if(line[13] == '!'): continue
            self.secArray.append(line[16])
	    self.phiArray.append(float(line[103:109]))
	    self.psiArray.append(float(line[109:115]))
	    i = int(line[6:10]) - 1
	    #as donor
	    j = int(line[41:45])
	    self.hbond[i][i+j] = float(line[46:50])
	    j = int(line[63:67])
	    self.hbond[i][i+j] = float(line[68:72])
	    #as acceptor
	    j = int(line[52:56])
	    self.hbond2[i][i+j] = float(line[57:61])
	    j = int(line[74:78])
	    self.hbond2[i][i+j] = float(line[79:83])

class PDB: 
    def __init__(self,fstr_pdb):
	fp = open(fstr_pdb,"r")
	self.name = fstr_pdb
	self._content = fp.readlines()
	self.coord = []
	for line in self._content:
	    if(line[:3] != "TER"):
		self.coord.append(float(line[30:38]))
		self.coord.append(float(line[38:46]))
		self.coord.append(float(line[46:54]))
	fp.close()

    def numberOfAtoms(self):
	n = 0
	for line in self._content:
	    if(line[:4] == 'ATOM'):
		n += 1
	return n

    def numberOfLinesInSnap(self,ifbox=0):
	"""determine how many lines the coordinations of a 
	conformation will occupay"""
	from math import ceil
	natoms = self.numberOfAtoms()
	n = ceil(natoms*3.0/10)
	size = natoms*3*8 + n
	if(ifbox):
	    n += 1
	    size += 25
	return int(n),int(size)
	
    def newpdb(self,crdstr,fstr="outpdb"):
	"""make a newpdb from a set of coordination (a snapshot from 
	trajectory file)"""
	outpdb = open(fstr,"w")
	i = 0
	if(type(crdstr[0]) == types.StringType):
	    for line in self._content:
		if(line[:4] != 'ATOM'):
		    outpdb.write(line)
		else:
		    outpdb.write(line[:30]+crdstr[i:i+24]+'\n')
		    i += 24
	elif(type(crdstr[0] == types.FloatType)):
	    for line in self._content:
		if(line[:4] != 'ATOM'):
		    outpdb.write(line)
		else:
		    outpdb.write("%s%8.3f%8.3f%8.3f\n" %(line[:30],crdstr[i],crdstr[i+1],crdstr[i+2]))
		    i += 3
	outpdb.close()

    def sas(self,fstr="",MSMS="~/MSMS/"):
	"""calculate solvent accessbile area and hydrophobic accessible area"""
	if fstr=="":
	    fstr = self.name
	system(MSMS+"pdb_to_xyzrn"+" "+fstr+" >temp.xyzrn")
	system(MSMS+"msms"+" "+" -noh -if temp.xyzrn -af temp >temp.log")

	fp = open("temp.area","r")
	lines = fp.readlines()
	
	sas = 0; hp = 0 # hydrophobic
	for i in range(1,len(lines)):
	    data = float(lines[i][16:22])
	    sas += data
	    if lines[i][24] == 'C': hp += data
	
	return sas, hp

def correlation(A,B):
    """calculate the correlation function between two sets of data.
       The definition of correlation function is based on equ(7.73) from 
       <<molecular modelling>>, Andrew Leach."""
    
    import Numeric as N
    from math import sqrt
    
    n = len(A)
    C = N.zeros((n,),N.Float)

    A = A - N.add.reduce(A)/n
    B = B - N.add.reduce(B)/n
    
    for i in range(n):
	C[i]=A[i]*B[i]

    nominator = N.add.reduce(C)
    denominator = sqrt((N.add.reduce(A**2)/n) * (N.add.reduce(B**2)/n))

    return nominator/denominator