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

import string
import coordcalc

def QuadCmp(a, b):
    return cmp(a.score, b.score)

class Quad:
    '''
    Scores and generates a quadruplet when given four residues.
    Rebuilds a quad when given 4 residues (already part of a quad)
    Deletes the quad (and all references)
    '''

    ### Private methods

    def __init__(self, parm, logger, level=3):
        self.parm = parm
        self.logger = logger
        self.level = level

    def __repr__(self):
        line = ''
        for res in self.res:
            res_str = res.get_parent().id
            res_str += '%4d' % (res.id[1])
            res_str += res.id[2]
            line += res_str + ' '
        for vtag in self.valid.keys():
            line += str(self.valid[vtag])+' '
        line += '%.4f' % (self.score)+'  '
        for type in [self.dist_score, self.bond_score, self.torsion_score]:
            for score in type:
                line += '%.4f' % score+' '
            line += ' '
        return line

    def _set_cutoff(self):
        level = self.level
        if self.orient == "parallel":
            if level == 1:
                self.cutoff = self.parm.syn_quad_cutoff
            elif level == 2:
                self.cutoff = self.parm.syn_quad_cutoff_2
        else:
            if level == 1:
                self.cutoff = self.parm.anti_quad_cutoff
            elif level == 2:
                self.cutoff = self.parm.anti_quad_cutoff_2

    def _failed_score(self):
        if hasattr(self, 'failed_param'):
            return self.failed_param
        if self.level == 3: return None
        return self.score > self.cutoff

    def _score_dist(self):
        '''
        Score quadruplets based on distance
        '''
        parm = self.parm
        self.dist_score = []
        self.dist = [[],[]]
        dist = [None] * 2
        if self.b1 and self.b2:
            dist[0] = coordcalc.distance \
                      (self.b1['CA'].coord, self.b2['CA'].coord)
        if self.c1 and self.c2:
            dist[1] = coordcalc.distance \
                      (self.c1['CA'].coord, self.c2['CA'].coord)
        for i in range(2):
            if not dist[i]:
                dist[i] = parm.syn_dist_mean
        for i in range(2):
            self.dist[i] += [dist[i]]
        ## check if a-b and c-d are within cutoffs
        for i in range(2):
            if dist[i] > parm.max_sheet_pair_cutoff:
                self.failed_param = 'dist'
                return
            if dist[i] < parm.min_sheet_pair_cutoff:
                self.failed_param = 'dist'
                return
        if self.orient == "parallel":
            weight = parm.syndistweight
            sigma = parm.syn_dist_sigma
            dist[0] -= parm.syn_dist_mean
            dist[1] -= parm.syn_dist_mean
        else:
            weight = parm.antidistweight
            sigma = parm.anti_dist_sigma
            if dist[0] > dist[1]:
                dist[0] -= parm.anti_dist_mean_hi
                dist[1] -= parm.anti_dist_mean_lo
            else:
                dist[0] -= parm.anti_dist_mean_lo
                dist[1] -= parm.anti_dist_mean_hi
        if weight == 0:
            return
        for i in range(2):
            dist[i] /= sigma
            self.dist[i] += [dist[i]]
        ## Check the range of the queries. At 6.0 sigma prob for a gaussian
        ## becomes zero and -ln() becomes undefined. Clip to 6 sigma (we use
        ## limits on distance above 3.5-7.5A)
        for i in range(2):
            if (dist[i] < -6.0) or (dist[i] > 6.0):
                dist[i] = 6.0
            dist[i] = self.parm.get_val('lpg', dist[i]);
            self.dist[i] += [dist[i]]
            dist[i] *= weight;
            self.score += dist[i];
            self.dist[i] += [dist[i]]
        self.dist_score = [dist[0], dist[1]]

    def _score_bond(self):
        '''
        Score quadruplets based on angle between lines joining bases
        and lines joining vertices (a1-c1 and b1-b2 etc.)
        '''
        parm = self.parm
        self.bond_score = []
        if self._failed_score(): return
        self.bond = [[],[],[],[]]
        bond = [None] * 4
        (a1, b1, c1, d1) = (self.a1, self.b1, self.c1, self.d1)
        (a2, b2, c2, d2) = (self.a2, self.b2, self.c2, self.d2)
        if a1 and c1 and b1 and b2:
            bond[0] = coordcalc.v_angle \
                      ((a1['CA'].coord - c1['CA'].coord),
                       (b1['CA'].coord - b2['CA'].coord), 90)
        if a2 and c2 and b2 and b1:
            bond[1] = coordcalc.v_angle \
                      ((a2['CA'].coord - c2['CA'].coord),
                       (b2['CA'].coord - b1['CA'].coord), 90)
        if b1 and d1 and c1 and c2:
            bond[2] = coordcalc.v_angle \
                      ((b1['CA'].coord - d1['CA'].coord),
                       (c1['CA'].coord - c2['CA'].coord), 90)
        if b2 and d2 and c2 and c1:
            bond[3] = coordcalc.v_angle \
                      ((b2['CA'].coord - d2['CA'].coord),
                       (c2['CA'].coord - c1['CA'].coord), 90)
        for i in range(4):
            if bond[i]: continue
            if self.orient == 'parallel':
                bond[i] = parm.syn_bond_mean
            else: bond[i] = parm.anti_bond_mean
        for i in range(4):
            self.bond[i] += [bond[i]]
        if self.orient == "parallel":
            weight = parm.synbondweight
            mean = parm.syn_bond_mean
            sigma = parm.syn_bond_sigma
        else:
            weight = parm.antibondweight
            mean = parm.anti_bond_mean
            sigma = parm.anti_bond_sigma
        if weight == 0: return
        for i in range(4):
            bond[i] = (bond[i] - mean)/sigma
            if not (-6.0 < bond[i] < 6.0):
                bond[i] = 6.0
            self.bond[i] += [bond[i]]
            if self.orient == "parallel":
                bond[i] = self.parm.get_val('sblp', bond[i])
            else:
                bond[i] = self.parm.get_val('ablp', bond[i])
            self.bond[i] += [bond[i]]
            bond[i] /= 2.0
            bond[i] *= weight
            self.score += bond[i]
            self.bond[i] += [bond[i]]
        self.bond_score = [bond[0], bond[1], bond[2], bond[3]]

    def _score_torsion_arrow(self):
        '''
        Score quadruplets based on torsion angles.
        '''
        parm = self.parm
        self.torsion_score = []
        if self._failed_score(): return
        self.tangle = [[],[]]
        angle = [None] * 2
        (a1, b1, c1, d1) = (self.a1, self.b1, self.c1, self.d1)
        (a2, b2, c2, d2) = (self.a2, self.b2, self.c2, self.d2)
        if b1 and a1 and c1 and a2 and c2 and b2:
            angle[0] = coordcalc.torsion \
                       (b1['CA'].coord,(a1['CA'].coord + c1['CA'].coord)/2,
                        (a2['CA'].coord + c2['CA'].coord)/2,b2['CA'].coord)
        if c1 and b1 and d1 and b2 and d2 and c2:
            angle[1] = coordcalc.torsion \
                       (c1['CA'].coord,(b1['CA'].coord+d1['CA'].coord)/2,
                        (b2['CA'].coord+d2['CA'].coord)/2,c2['CA'].coord)
        for i in range(2):
            if angle[i]: continue
            if self.orient == 'parallel':
                angle[i] = parm.syn_torsion_arrow_mean
            else: angle[i] = parm.anti_torsion_arrow_mean
        for i in range(2):
            self.tangle[i] += [angle[i]]
        if self.orient == "parallel":
            weight = parm.syntorsionarrowweight
            mean = parm.syn_torsion_arrow_mean
            sigma = parm.syn_torsion_arrow_sigma
        else:
            weight = parm.antitorsionarrowweight
            mean = parm.anti_torsion_arrow_mean
            sigma = parm.anti_torsion_arrow_sigma
        if weight == 0: return
        for i in range(2):
            angle[i] = (angle[i] - mean) / sigma;
            if not (-6.0 < angle[i] < 6.0):
                angle[i] = 6.0
            self.tangle[i] += [angle[i]]
            if self.orient == "parallel":
                angle[i] = self.parm.get_val('stlp', angle[i])
            else:
                angle[i] = self.parm.get_val('atlp', angle[i])
            self.tangle[i] += [angle[i]]
            angle[i] *= weight;
            self.score += angle[i];
            self.tangle[i] += [angle[i]]
        self.torsion_score = [angle[0], angle[1]]

    def _validate(self):
        '''
        Check if quadruplet can be used for linking to form ladders.
        Simply test with a cutoff score (from parameters file).
        '''
        if not (self.dist_score and self.bond_score and self.torsion_score):
            return
        self.valid['q3'] = 1
        if self.orient == "parallel":
            self.valid['q1'] = (self.score < self.parm.syn_quad_cutoff)
            self.valid['q2'] = (self.score < self.parm.syn_quad_cutoff_2)
        else:
            self.valid['q1'] = (self.score < self.parm.anti_quad_cutoff)
            self.valid['q2'] = (self.score < self.parm.anti_quad_cutoff_2)

    def _tracequad(self):
        lines = []
        if hasattr(self, 'dist'):
            for i in range(2):
                lines += [str(self.dist[i])+'\n']
        if hasattr(self, 'bond'):
            for i in range(4):
                lines += [str(self.bond[i])+'\n']
        if hasattr(self, 'tangle'):
            for i in range(2):
                lines += [str(self.tangle[i])+'\n']
        return lines

    def _pairup(self, res1, res2, type):
        if res1 in res2.pairs['all']: return
        if res2 in res1.pairs['all']: return
        if res1 in res2.blg: return
        if res2 in res1.blg: return
        if hasattr(res1, 'hpin') and res1.hpin == res2: return
        if hasattr(res2, 'hpin') and res2.hpin == res1: return
        res1.pairs['all'].append(res2)
        res2.pairs['all'].append(res1)
        res1.pairs[type].append(res2)
        res2.pairs[type].append(res1)

    ### Public methods

    def log(self, handle=None):
        line = 'Quad '
        line += self.__repr__()
        #line += '\n'
        self.logger.log('quad', line, handle)
        lines = self._tracequad()
        self.logger.log('quadtrace', lines, handle)

    def generate(self,
                  a1, b1, c1, d1,
                  a2, b2, c2, d2, orient=None):
        '''
        This generates a quadruplet and scores it.
        quad is:
         ||  a1      a2  /\  (anti-parallel)
         ||  b1  --  b2  ||
         ||  c1  --  c2  ||
         \/  d1      d2  \/  (parallel)
        '''
        (self.a1, self.b1, self.c1, self.d1) = (a1, b1, c1, d1)
        (self.a2, self.b2, self.c2, self.d2) = (a2, b2, c2, d2)
        if orient: self.orient = orient
        self.res = [b1,b2,c1,c2]
        self.pair = {b1:b2, b2:b1, c1:c2, c2:c1}
        self.link = {'b':None, 'c':None}
        self.score = 0
        self.valid = {}
        for q in ['q1','q2','q3']:
            self.valid[q] = None
        self._set_cutoff()
        self._score_dist()
        self._score_bond()
        self._score_torsion_arrow()
        self._validate()
        
    def build(self, resi, resj, orientation):
        self.orient = orientation
        if orientation == 'parallel':
            self.generate(resi.prev[1], resi, resi.next[1], resi.next[2],
                           resj.prev[1], resj, resj.next[1], resj.next[2])
        elif orientation == 'antiparallel':
            self.generate(resi.prev[1], resi, resi.next[1], resi.next[2],
                           resj.next[2], resj.next[1], resj, resj.prev[1])
        #self.log('quad.build')

    def is_similar_to(self, other):
        '''
        quads are similar if 3 or more residues are common
        quads are similar if edges are within one residue
        '''
        if self == other:
            return 'same'
        ctr = 0
        for i in self.res:
            for j in other.res:
                if i == j: ctr += 1
        if ctr > 2:
            return 'similar'
        if self.b1 == other.b1 and \
               (self.b2 in [other.b2.prev[1], other.b2, other.b2.next[1]]):
            return 'similar'
        if self.b2 == other.b2 and \
               (self.b1 in [other.b1.prev[1], other.b1, other.b1.next[1]]):
            return 'similar'
        return None

    def checkandfit(self, _quad):
        '''
        check if a free edge of self fits a free edge of _quad
        '''
        if self.orient != _quad.orient: return

        if self.link['b'] and self.link['c']:
            return
        if _quad.link['b'] and _quad.link['c']:
            return

        if not (self.link['c'] or _quad.link['b']):            
            if (self.c1 == _quad.b1 and \
                (self.c2 in [_quad.b2.prev[1], _quad.b2, _quad.b2.next[1]])) \
                or \
                (self.c2 == _quad.b2 and \
                 (self.c1 in [_quad.b1.prev[1], _quad.b1, _quad.b1.next[1]])):

                self.link['c'] = _quad
                _quad.link['b'] = self
                if _quad.b1 != self.c1:
                    _quad.b1.blg.append(_quad.b2)
                if _quad.b2 != self.c2:
                    _quad.b2.blg.append(_quad.b1)
                return
        if not (self.link['b'] or _quad.link['c']):
            if (self.b1 == _quad.c1 and \
                (self.b2 in [_quad.c2.prev[1], _quad.c2, _quad.c2.next[1]])) \
                or \
                (self.b2 == _quad.c2 and \
                 (self.b1 in [_quad.c1.prev[1], _quad.c1, _quad.c1.next[1]])):
                self.link['b'] = _quad
                _quad.link['c'] = self
                if _quad.c1 != self.b1:
                    _quad.c1.blg.append(_quad.c2)
                if _quad.c2 != self.b2:
                    _quad.c2.blg.append(_quad.c1)

    def isolate(self):
        if self.link['b']:
            self.link['b'].link['c'] = None
        if self.link['c']:
            self.link['c'].link['b'] = None
        for res in self.res:
            if hasattr(res, 'blg'):
                if self.pair[res] in res.blg:
                    res.blg.remove(self.pair[res])

    def pairup(self, type):
        self._pairup(self.b1, self.b2, type)
        self._pairup(self.c1, self.c2, type)
        for res in self.res:
            setattr(res, type, 1)
