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

from CoordCalc import v_angle

class HelixbrkStat:
    def __init__(self, logger):
        # recursive
        self.statorig = '../dat/stat/Hrmsd-orig'
        self.statbrok = '../dat/stat/Hrmsd-brok'
        self.statangle = '../dat/stat/Hangle'
        self.statanglebin = '../dat/stat/Hanglebin'
        self.logger = logger
        self.rmsdorig = {}
        self.rmsdbrok = {}
        self.angle = []
        self.anglebin = {}
        for i in range(100):
            self.rmsdorig[i] = []
            self.rmsdbrok[i] = []
        for i in range(0,100,5):
            self.anglebin[i] = []
        # incremental angles
        self.conseq = []
        self.ov0 = []
        self.ov1 = []
        self.ov2 = []

    # by recursion break (minimising Rmax of broken helices)
    def stathelix_recursive(self, helix):
        length = []
        rmsd = []
        angle = []
        for brks in range(4):
            if (brks > (len(helix)/5)):
                continue
            helices = helix.getBrokenHelices(brks)
            le = []
            r = []
            a = []
            for h in helices:
                le.append(len(h))
                r_ = h.getRMSD()
                r.append(r_)
                if (brks == 0):
                    self.rmsdorig[len(h)].append(r_)
                else:
                    self.rmsdbrok[len(h)].append(r_)
            for (i,h) in enumerate(helices[:-1]):
                ang=h.getAngle(helices[i+1])
                a.append(ang)
                self.angle.append(ang)
            length.append(le)
            rmsd.append(r)
            angle.append(a)
            if (a and ((max(a) - min(a)) < 5.0)):
                ang = sum(a)/len(a)
                self.anglebin[5 * int(ang/5.0)] += r
        helix.log('stathelix')
        self.logger.log('stat', 'stathelix length\t'+str(length))
        self.logger.log('stat', 'stathelix rmsd\t'+str(rmsd))
        self.logger.log('stat', 'stathelix angle\t'+str(angle))

    def dumphelixstat_recursive(self):
        lines = []
        for i in range(100):
            if self.rmsdorig[i]:
                for rmsd in self.rmsdorig[i]:
                    line = [str(i)+'\t'+str(rmsd)+'\n']
                    lines += line
        fh = file(self.statorig, "a", 1)
        fh.writelines(lines)
        fh.close()
        lines = []
        for i in range(100):
            if self.rmsdbrok[i]:
                for rmsd in self.rmsdbrok[i]:
                    line = [str(i)+'\t'+str(rmsd)+'\n']
                    lines += line
        fh = file(self.statbrok, "a", 1)
        fh.writelines(lines)
        fh.close()
        lines = []
        for i in self.angle:
            line = [str(i)+'\n']
            lines += line
        fh = file(self.statangle, "a", 1)
        fh.writelines(lines)
        fh.close()
        lines = []
        for i in range(0,100,5):
            if self.anglebin[i]:
                for rmsd in self.anglebin[i]:
                    line = [str(i)+'\t'+str(rmsd)+'\n']
                    lines += line
        fh = file(self.statanglebin, "a", 1)
        fh.writelines(lines)
        fh.close()

    def stathelix_incremental(self, helix):
        pass

    def dumphelixstat_incremental(self, helix):
        pass

class StrandbrkStat:
    def __init__(self, logger):
        self.stati2 = '../dat/stat/Eangle_i2'
        self.stati4 = '../dat/stat/Eangle_i4pseudo'
        self.stati2i4 = '../dat/stat/Eangle_i2_i4'
        self.logger = logger
        self.anglei2 = []
        self.anglei4 = []
        self.angledicti2 = {}
        self.angledicti4 = {}

    def statstrand_i2_i4(self, element):
        if len(element) < 6: return # minimum 6 for bulge
        i2angles = []
        for (i,res) in zip(range(len(element)),element)[2:-3]:
            i2Angle=v_angle((element[i-2]['CA'].coord-res['CA'].coord),
                            (res['CA'].coord-element[i+2]['CA'].coord), 180)
            i2angles.append(i2Angle)
        self.anglei2 += i2angles


        if len(element) < 6: return # minimum 6 for bulge
        i4angles = []
        for (i,res) in zip(range(len(element)),element)[2:-3]:
            vCAcoord = (element[i-1]['CA'].coord + \
                        res['CA'].coord + \
                        element[i+1]['CA'].coord + \
                        element[i+2]['CA'].coord) / 4
            i4pseudoAngle=v_angle((element[i-2]['CA'].coord-vCAcoord),
                                  (vCAcoord-element[i+3]['CA'].coord), 180)
            i4angles.append(i4pseudoAngle)
        self.anglei4 += i4angles

        if len(element) < 6: return
        for (i, i2) in enumerate(i2angles):
            for (j, i4) in enumerate(i4angles):
                ky = j-i
                if (abs(ky) > 6): continue
                if not self.angledicti2.has_key(ky):
                    self.angledicti2[ky]=[]
                if not self.angledicti4.has_key(ky):
                    self.angledicti4[ky]=[]
                self.angledicti2[ky].append(i2)
                self.angledicti4[ky].append(i4)

    def dumpangledict(self):
        stati2i4 = self.stati2i4
        for ky in self.angledicti2.keys():
            fname = stati2i4 + '_' + str(ky)
            lines = []
            for (i, a) in enumerate(self.angledicti2[ky]):
                lines.append(str(a)+'\t'+str(self.angledicti4[ky][i])+'\n')
            fh = file(fname, 'a', 1)
            fh.writelines(lines)
            fh.close()
    
    def dumpstrandstat_i2_i4(self):
        lines = []
        for a in self.anglei2:
            lines.append(str(a)+'\n')
        fh = file(self.stati2, 'a', 1)
        fh.writelines(lines)
        fh.close()
        lines = []
        for a in self.anglei4:
            lines.append(str(a)+'\n')
        fh = file(self.stati4, 'a', 1)
        fh.writelines(lines)
        fh.close()
        self.dumpangledict()

class ElemDistOvlapStat:
    def __init__(self, logger):
        self.logger = logger
        self.statelemHHdist = '../dat/stat/elemHHdist'
        self.statelemHEdist = '../dat/stat/elemHEdist'
        self.statelemEEdist = '../dat/stat/elemEEdist'
        self.statelemHHovlap = '../dat/stat/elemHHovlap'
        self.statelemHEovlap = '../dat/stat/elemHEovlap'
        self.statelemEEovlap = '../dat/stat/elemEEovlap'
        self.HHdists = []
        self.HEdists = []
        self.EEdists = []
        self.HHovlaps = []
        self.HEovlaps = []
        self.EEovlaps = []

    def statElem(self, elem1, elem2):
        dist = elem1.getDist(elem2)
        ovlap = elem1.getOverlap(elem2)
        if ((elem1.type == 'H') and (elem2.type == 'H')):
            self.HHdists.append(dist)
            self.HHovlaps.append(ovlap)
        elif ((elem1.type == 'E') and (elem2.type == 'E')):
            self.EEdists.append(dist)
            self.EEovlaps.append(ovlap)
        else:
            self.HEdists.append(dist)
            self.HEovlaps.append(ovlap)

    def dumpelemstat(self):
        ary = [
            [self.statelemHHdist, self.HHdists],
            [self.statelemHEdist, self.HEdists],
            [self.statelemEEdist, self.EEdists],
            [self.statelemHHovlap, self.HHovlaps],
            [self.statelemHEovlap, self.HEovlaps],
            [self.statelemEEovlap, self.EEovlaps]]
        for (fn, values) in ary:
            lines = []
            for a in values:
                lines.append(str(a)+'\n')
            fh = file(fn, 'a', 1)
            fh.writelines(lines)
            fh.close()
            
        
