# $CVSHeader: pal8e/src/PhiPsi.py,v 1.2 2005/02/15 01:10:26 indraneel Rel $

# create data for Ramachandran plot from helix and strand defined by PSSD

from Log import LogFH
import os

from PSSDParser import PSSD

class Parm:
    gridsize = 2

class StatKeeper:
    def __init__(self, parm):
        self.parm = parm
        self.helixData = []
        self.strandData = []

    def _get_grid(self, typ):
        parm = self.parm
        gridData = {}
        for phi in range(-180/parm.gridsize, 180/parm.gridsize):
            Phi = (phi*parm.gridsize)+(parm.gridsize/2.0)
            gridData[Phi] = {}
            for psi in range(-180/parm.gridsize, 180/parm.gridsize):
                Psi = (psi*parm.gridsize)+(parm.gridsize/2.0)
                gridData[Phi][Psi] = {}
                gridData[Phi][Psi][typ] = 0
        return gridData

    def addData(self, data, typ):
        if typ == 'H':
            self.helixData += data
        elif typ == 'E':
            self.strandData += data

    def getData(self, typ):
        if typ == 'H':
            return self.helixData
        elif typ == 'E':
            return self.strandData

    def getGridData(self, typ):
        parm = self.parm
        gridData = self._get_grid(typ)
        if typ == 'H':
            data = self.helixData
        elif typ == 'E':
            data = self.strandData
        for [phi, psi] in data:
            if phi == 360.0: continue
            if psi == 360.0: continue
            if phi == 180.0: phi = 179.99
            if psi == 180.0: psi = 179.99
            Phi = int(phi/parm.gridsize) * parm.gridsize + \
                  (parm.gridsize/2.0)
            Psi = int(psi/parm.gridsize) * parm.gridsize + \
                  (parm.gridsize/2.0)
            gridData[Phi][Psi][typ] += 1
        return gridData

    def getNormGridData(self, typ):
        gridData = self.getGridData(typ)
        total = 0
        for phi in gridData.keys():
            for psi in gridData[phi].keys():
                total += gridData[phi][psi][typ]
        total = float(total)
        for phi in gridData.keys():
            for psi in gridData[phi].keys():
                gridData[phi][psi][typ] /= total
        return gridData

    def getGridDataList(self, typ):
        gridData = self.getGridData(typ)

    def getNormGridDataList(self, typ):
        gridData = self.getNormGridData(typ)
        data = []
        for phi in gridData.keys():
            for psi in gridData[phi].keys():
                data.append([phi, psi, gridData[phi][psi][typ]])
        return data

class PhiPsi:
    def __init__(self, logfile, parm):
        self.statkeeper = StatKeeper(parm)
        self.logger = LogFH(logfile)
        
    def runfile(self, fname):
        #self.logger.log(fname)
        pssd = PSSD(self.logger)
        pssd.store(fname)
        fid = os.path.basename(fname)
        for elem in pssd.elements:
            for typ in ['H', 'E']:
                if elem.type == typ:
                    self.statkeeper.addData(elem.getPhiPsi(handle=fid), typ)

    def runall(self, pssddir, outfname):
        for (i, f) in enumerate(os.listdir(pssddir)):
            fname = pssddir+f
            self.runfile(fname)
            #if i > 50: break
        for typ in ['H', 'E']:
            data = self.statkeeper.getData(typ)
            lines = []
            for d in data:
                lines.append("%.1f" % (d[0])+'\t'+"%.1f" % (d[1])+'\n')
            ofh = file(outfname+typ, 'w')
            ofh.writelines(lines)
            #griddata = self.statkeeper.getNormGridDataList(typ)
            #lines = []
            #for d in griddata:
            #    lines.append(str(d[0])+'\t'+str(d[1])+'\t'+str(d[2])+'\n')
            #ofh = file(outfname+typ+'grid', 'w')
            #ofh.writelines(lines)
            
if __name__ == '__main__':
    pssddir = '../../dat/cull100.pssd/'
    logfile = 'x'
    outfile = 'y'
    parm = Parm()
    pp = PhiPsi(logfile, parm)
    pp.runall(pssddir, outfile)
    
