from __future__ import division

from phaser import ncs

from scitbx.array_family import flex

import re
import difflib
from collections import deque
import math
import operator

def localize(sequence, width, padding):
    """
    Groups values and pads them to a fixed length
    """
        
    assert 0 <= width
    padding = [ padding ] * width
        
    padded = padding + sequence + padding
    d = deque(
        iterable = padded[ : 2 * width ],
        maxlen = 2 * width + 1
        )
    
    for v in padded[ 2 * width : ]:
        d.append( v )
        yield d


def split(sequence, consecutivity):
    """
    Breaks up sequences where elements are not consecutive according to a
    property
    """
                
    while sequence:
        iter = enumerate( sequence )
        ( index, left ) = iter.next()
        
        for ( index, right ) in iter:
            if not consecutivity( left = left, right = right ):
                index -= 1
                break
            
            left = right
        
        divisor = index + 1
        leader = sequence[ : divisor ]
        sequence = sequence[ divisor : ]
        yield leader
    
    
class StructuralConsecutivity(object):
    """
    Two residue_groups are consecutive if:
        distance( pre of left, post of right ) <= distance_cutoff
    NB. If either ( pre of left, post of right ) is missing, the two
        are not consecutive
    If a residue_group is disordered, only 'A'-'A' and ' '-'A' type bonds are
        considered, the rest is assumed to be correct 
    """
    
    def __init__(self, pre_bond_atom, post_bond_atom, distance_cutoff):
        
        self.pre_bond_atom = pre_bond_atom
        self.post_bond_atom = post_bond_atom
        self.distance_cutoff = distance_cutoff
        
        
    def __call__(self, left, right):
    
        pre_atom_with = self.candidates( rg = left, name = self.pre_bond_atom )
        post_atom_with = self.candidates( rg = right, name = self.post_bond_atom )
        
        if not pre_atom_with or not post_atom_with:
            return False
        
        if "" in pre_atom_with:
            pre_atom = pre_atom_with[ "" ]
            post_atom = post_atom_with[ min( post_atom_with ) ]
            
        elif "" in post_atom_with:
            post_atom = post_atom_with[ "" ]
            pre_atom = pre_atom_with[ min( pre_atom_with ) ]
            
        else:
            common = sorted( [ a for a in pre_atom_with if a in post_atom_with ] )
            
            if not common:
                return False
            
            pre_atom = pre_atom_with[ common[0] ]
            post_atom = post_atom_with[ common[0] ]
        
        return pre_atom.distance( post_atom ) <= self.distance_cutoff
        
        
    @staticmethod    
    def candidates(rg, name):
        
        return dict(
            [ ( a.parent().altloc, a ) for a in rg.atoms() if a.name == name ]
            )
        
        
def map_residues_to_sequence(
    sequence,
    residues,
    consecutivity,
    one_letter_for,
    min_length,
    gap_char
    ):
    """
    Aligns residue list with sequence 
    """ 
        
    ( indices, gapless ) = zip(
        *[ p for p in enumerate( sequence ) if p[1] != gap_char ]
        )
    gapless = "".join( gapless )
    mapping = [ None ] * len( sequence )
    current = 0
    
    for fragment in split( sequence = residues, consecutivity = consecutivity ):
        if len( fragment ) < min_length:
            continue
        
        # A regular expression is used to align fragment to sequence
        fragment_seq = "".join(
            [ one_letter_for.get( rg.atom_groups()[0].resname, "." )
            for rg in fragment ]
            )
        match = re.search( fragment_seq, gapless[ current : ] )
        
        if not match:
            raise RuntimeError, (
                "Full sequence = %s, Fragment sequence = %s, (start = %s)" % (
                    gapless,
                    fragment_seq,
                    current
                    )
                ) 
        
        start = match.start( 0 ) + current
        
        for ( i, residue_group ) in enumerate( fragment, start = start ):
            assert i < len( indices )
            assert indices[ i ] < len( mapping )
            mapping[ indices[ i ] ] = residue_group
            
        current = i
        
    return mapping


def sequence_overlap(left, right):
    
    lcs = ncs.LCSTree()
    lcs.insert_word( left )
    lcs.insert_word( right )
    return lcs.get_longest_common_substring()


def sequence_shift(left, right, lcs):
    
    start_1 = left.find( lcs )
    assert start_1 != -1
    start_2 = right.find( lcs )
    assert start_2 != -1
    
    forw = start_1 - start_2
    midw = len( left ) - forw
    backw = midw - len( right )
    
    return ( forw, midw, backw )


def sequence_stich(seq_extend, seq_final):
    
    ( forw, midw, backw ) = sequence_shift(
        left = seq_final,
        right = seq_extend,
        lcs = sequence_overlap( left = seq_final, right = seq_extend ),
        )
    start = max( 0, forw )
    length = len( seq_extend )
    return seq_final[ : start ] + seq_extend + seq_final[ start + length : ]


def alignment_stich(sequence, alignment, index):
    
    assert index < alignment.multiplicity()
    sequences = alignment.sequence_strings()
    sequences[ index ] = sequence_stich(
        seq_extend = sequences[ index ],
        seq_final = sequence.sequence,
        )
    alignment.extend( sequences = sequences )
        
    
def sequence_match(left, right, gap_char):
    
    return (
        gap_char * -forw + left + gap_char * -backw,
        gap_char * forw + right + gap_char * backw,
        )


def align_residues_to_sequence(
    residues,
    sequence,
    consecutivity,
    one_letter_for,
    unknown,
    min_hssp_length = 6,
    min_fragment_length = 4,
    min_fraction = 0.4,
    ):
    """
    Aligns residue list with sequence
    """
        
    alignment = []
    remaining = sequence
    
    for fragment in split( sequence = residues, consecutivity = consecutivity ):
        if len( fragment ) < min_fragment_length:
            continue
        
        fragment_seq = one_letter_sequence(
            rgs = fragment,
            one_letter_for = one_letter_for,
            unknown = unknown
            )
        l = len( fragment )
        lcs = ncs.LCSTree()
        lcs.insert_word( fragment_seq )
        lcs.insert_word( remaining )
        ss = lcs.get_longest_common_substring()
        
        if len( ss ) < min_hssp_length and len( ss ) < l:
            continue
        
        start_1 = fragment_seq.find( ss )
        assert start_1 != -1
        start_2 = remaining.find( ss )
        assert start_2 != -1
        forw = start_1 - start_2
        midw = l - forw
        backw = midw - len( remaining )
        f_ali = [ None ] * -forw + zip( fragment, fragment_seq )
        s_ali = [ None ] * forw + list( remaining[ : midw ] ) + [ None ] * backw
        
        matching = len( [ p for ( p, s ) in zip( f_ali, s_ali )
            if p is not None and p[1] == s ] )
                
        if matching < min( l, len( remaining ) ) * min_fraction:
            continue
        
        alignment.extend(
            p[0] if p is not None else None for ( p, s ) in zip( f_ali, s_ali )
            if s is not None
            )
        remaining = remaining[ midw : ]
        
        if not remaining:
            break
        
    alignment.extend( [ None ] * len( remaining ) )
    
    assert len( alignment ) == len( sequence )
    
    return alignment


def align_chain_to_sequence(
    chain,
    sequence,
    consecutivity,
    one_letter_for,
    unknown,
    min_hssp_length,
    min_fraction,
    gap_char
    ):
    """
    Aligns residue list with sequence 
    """
        
    first = []
    second = []
    remaining = sequence
    
    for fragment in split( sequence = chain.residue_groups(), consecutivity = consecutivity ):
        fragment_seq = one_letter_sequence(
            rgs = fragment,
            one_letter_for = one_letter_for,
            unknown = unknown
            )
        l = len( fragment )
        lcs = ncs.LCSTree()
        lcs.insert_word( fragment_seq )
        lcs.insert_word( remaining )
        ss = lcs.get_longest_common_substring()
        
        if len( ss ) < min_hssp_length and len( ss ) < l:
            first.append( fragment )
            second.append( gap_char * l )
            continue
        
        start_1 = fragment_seq.find( ss )
        assert start_1 != -1
        start_2 = remaining.find( ss )
        assert start_2 != -1
        forw = start_1 - start_2
        midw = l - forw
        backw = midw - len( remaining )
        f_ali = [ None ] * -forw + fragment
        s_ali = gap_char * forw + remaining[ : midw ] + gap_char * backw
        
        matching = 0
        
        for ( f, s ) in zip( f_ali, s_ali ):
            if not f:
                continue
            
            if one_letter_for.get( f.atom_groups()[0].resname, unknown ) == s:
                matching += 1
                
        if matching < min( l, len( remaining ) ) * min_fraction:
            first.append( fragment )
            second.append( gap_char * l )
            continue
        
        first.append( f_ali )
        second.append( s_ali )
            
        remaining = remaining[ midw : ]
        
    first.append( [ None ] * len( remaining ) )
    second.append( remaining )
        
    return ( reduce( operator.add, first, [] ), reduce( operator.add, second, "" ) )


def match_chain_to_alignment(chain, mtype, sequence, gap, min_hssp_length, min_matching_fraction):
    
    gapless = "".join( c for c in sequence if c != gap )
    ( a_rgs, a_seq ) = align_chain_to_sequence(
        chain = chain,
        sequence = gapless,
        consecutivity = mtype.consecutivity,
        one_letter_for = mtype.one_letter_for(),
        unknown = mtype.unknown_one_letter,
        min_hssp_length = min_hssp_length,
        min_fraction = min_matching_fraction,
        gap_char = gap
        )
    
    assert len( a_rgs ) == len( a_seq )
    assert "".join( c for c in a_seq if c != gap ) == gapless
    mapping = [ None ] * len( sequence )
    indices = [ i for ( i, c ) in enumerate( sequence ) if c != gap ]
    
    for ( i, ( rg, c ) ) in enumerate( [ p for p in zip( a_rgs, a_seq ) if p[1] != gap ] ):
        assert i < len( indices )
        assert indices[ i ] < len( mapping )
        mapping[ indices[ i ] ] = rg
        
    return mapping


def mismatching_sequence_alignment(
    residues,
    sequence,
    consecutivity,
    one_letter_for,
    min_hssp_length,
    unknown,
    gap_char
    ):
    """
    Aligns residue list with sequence 
    """
        
    first = []
    second = []
    remaining = sequence
    
    for fragment in split( sequence = residues, consecutivity = consecutivity ):
        fragment_seq = "".join(
            [ one_letter_for.get( rg.atom_groups()[0].resname, unknown )
            for rg in fragment ]
            )
        l = len( fragment_seq )
        lcs = ncs.LCSTree()
        lcs.insert_word( fragment_seq )
        lcs.insert_word( remaining )
        ss = lcs.get_longest_common_substring()
        
        if len( ss ) < min_hssp_length and len( ss ) < l:
            first.append( fragment_seq )
            second.append( gap_char * l )
            continue
        
        start_1 = fragment_seq.find( ss )
        assert start_1 != -1
        start_2 = remaining.find( ss )
        assert start_2 != -1
        forw = start_1 - start_2
        midw = l - forw
        backw = midw - len( remaining )
        first.append( gap_char * -forw + fragment_seq )
        second.append( gap_char * forw + remaining[ : midw ] + gap_char * backw )
        remaining = remaining[ midw : ]
        
    return ( reduce( operator.add, first ), reduce( operator.add, second ) )


def map_to_sequences(
    sequences,
    residues,
    consecutivity,
    one_letter_for,
    min_length,
    gap_char
    ):
    
    mappings = []
        
    for seq in sequences:
        try:
            resi_map = map_residues_to_sequence(
                sequence = seq,
                residues = residues,
                consecutivity = consecutivity,
                one_letter_for = one_letter_for,
                min_length = min_length,
                gap_char = gap_char
                )
        
        except RuntimeError, e:
            resi_map = None
        
        mappings.append( resi_map )
        
    return mappings


def map_to_alignment(alignment, residues, consecutivity, one_letter_for, min_length):
    
    for ( index, seq ) in enumerate( alignment.alignments ):
        try:
            resi_map = map_residues_to_sequence(
                sequence = seq,
                residues = residues,
                consecutivity = consecutivity,
                one_letter_for = one_letter_for,
                min_length = min_length,
                gap_char = alignment.gap
                )
        
        except RuntimeError:
            continue
        
        return ( index, resi_map )
    
    return ( -1, [] )


def residue_group_difference(left, right):
    
    substrahend = set( [ r for r in right if r ] )
    return [ rg for rg in left if rg not in substrahend ]


def chain_residue_groups_difference(chain, residue_groups):
    
    substr = set( [ rg.memory_id() for rg in residue_groups  if rg ] )
    return [ rg for rg in chain.residue_groups() if rg.memory_id() not in substr ]


def three_letter_sequence(rgs):
    
    return [ rg.atom_groups()[0].resname for rg in rgs if rg and rg.atom_groups() ]


def one_letter_sequence(rgs, one_letter_for, unknown):
    
    return "".join( [ one_letter_for.get( s, unknown )
        for s in three_letter_sequence( rgs = rgs ) ] )

   
class NumberingFrame(object):
    """
    Calculates frame shift wrt master residue group sequence
    """
    
    def __init__(self, residues1, residues2, consecutivity):
        
        if not residues1 or not residues2:
            self.shift = 0
            self.overlap = 0
        
        else:
            matcher = difflib.SequenceMatcher()
            
            fragments1 = list(
                split( sequence = residues1, consecutivity = consecutivity )
                )
            sequences1 = [ self.to_sequence( residues = f ) for f in fragments1 ]
            matches = []
            
            # SequenceMatcher caches seq2
            for frag2 in split( sequence = residues2, consecutivity = consecutivity ):
                seq2 = self.to_sequence( residues = frag2 )
                matcher.set_seq2( seq2 )
                
                for ( frag1, seq1 ) in zip( fragments1, sequences1 ):
                    matcher.set_seq1( seq1 )
                    ( a, b, length ) = matcher.find_longest_match( 0, len( seq1 ), 0, len( seq2 ) )
                    
                    matches.append(
                        ( frag1[ a ], frag2[ b ], length )
                        )
                    
            ( r1, r2, length ) = max( matches, key = lambda s: s[2] )
            
            self.shift = r2.resseq_as_int() - r1.resseq_as_int()
            self.overlap = length
        
    
    @staticmethod
    def to_sequence(residues):
        
        return [ rg.atom_groups()[0].resname for rg in residues ]
    

class VanDerWaalsRadius(object):
    """
    Supplies van der Waals-radii of given type
    """
    
    def __init__(
        self,
        root,
        radius_for_mainchain_atom,
        sidechain_data_for,
        unknown_radius
        ):
        
        atoms = root.atoms()
        self.values = flex.double()
        self.values.reserve( len( atoms ) )
        
        for atom in atoms:
            if atom.name in radius_for_mainchain_atom:
                self.values.append( radius_for_mainchain_atom[ atom.name ] )
            
            else:
                radius_for = sidechain_data_for.get( atom.parent().resname, {} )
                self.values.append( radius_for.get( atom.name, unknown_radius ) )
    
    
    @classmethod
    def from_mmt(cls, root, mmt):
        
        return cls(
            root = root,
            radius_for_mainchain_atom = dict(
                zip( mmt.mainchain_atom_names, mmt.mainchain_atom_radii )
                ),
            sidechain_data_for = dict(
                zip( mmt.three_letter_codes, mmt.sidechain_atom_radii_dicts )
                ),
            unknown_radius = mmt.unknown_atom_radius
            )
        
        
def distance_sqs_from(centre, coords):
    
    diffs = coords - centre
    return diffs.dot( diffs )
    
    
class AccessibleSurfaceArea(object):
    """
    Calculates accessible surface area values
    """
    
    def __init__(self, root, vdw_radii, probe = 1.4, precision = 960):
        
        assert 0 <= probe, "Negative probe radius: %s" % probe
        assert 0 <= precision, "Negative precision: %s" % precision
        
        atoms = list( root.atoms() ) # enable identity-based comparison
        coords = root.atoms().extract_xyz() 
        probing_radii_sqs = self.probing_radii_squares(
            radii = vdw_radii.values,
            probe = probe
            )
        altlocs = [ a.parent().altloc for a in atoms ] 
        sphere_points = self.points_on_sphere( count = precision )
        accessible_counts = flex.double()
        accessible_counts.reserve( len( atoms ) )
        max_distances = vdw_radii.values + 2.0 * probe
        
        for ( atom, radius, maxdist ) in zip( atoms, vdw_radii.values, max_distances ):
            distances = flex.sqrt(
                distance_sqs_from( centre = atom.xyz, coords = coords )
                )
            cutoff = vdw_radii.values + maxdist
            selected = self.atoms_to_consider(
                centre = atom,
                selection = distances < cutoff,
                atoms = atoms,
                probing_radii_sqs = probing_radii_sqs,
                altlocs = altlocs
                )
            trial_points = (
                ( radius + probe ) * sphere_points + atom.xyz
                )
            accessible_counts.append(
                self.count_accessible_points(
                    points = trial_points,
                    considered_neighbour_data = selected
                    )
                )
        
        self.values = ( 4.0 * math.pi / precision
            * accessible_counts * probing_radii_sqs )
        
        
    @staticmethod
    def count_accessible_points(points, considered_neighbour_data):
            
        for ( trial, probing_radius_sq ) in considered_neighbour_data:
            distance_sqs = distance_sqs_from( centre = trial, coords = points )
            points = points.select( probing_radius_sq <= distance_sqs ) 
                
        return len( points )
        
        
    @staticmethod
    def probing_radii_squares(radii, probe):
        
        probing_radii = radii + probe
        return probing_radii * probing_radii
    
    
    @staticmethod
    def atoms_to_consider(
        centre,
        selection,
        atoms,
        probing_radii_sqs,
        altlocs
        ):
        
        centre_atom_altloc = centre.parent().altloc 
        base_selector = lambda a, s, aloc: a != centre and s
            
        if centre_atom_altloc:
            selector = lambda a, s, aloc: (
                ( not aloc or aloc == centre_atom_altloc )
                and base_selector( a, s, aloc )  
                )
            
        else:
            selector = base_selector
            
        return [ ( a.xyz, prs ) for ( a, s, aloc, prs )
                in zip( atoms, selection, altlocs, probing_radii_sqs )
                if selector( a, s, aloc ) ]
    
    
    @staticmethod
    def points_on_sphere(count):
        "Generates quasi equidistant points on a sphere surface"
    
        # Calculate values for quick access...
        increment = math.pi * ( 3.0 - math.sqrt( 5.0 ) )
        offset = 2.0 / count
    
        # Calculate points...
        indices = flex.double_range( count )
        ys = offset * indices - 1.0 + ( offset / 2.0 )
        rs = flex.sqrt( 1.0 - ys * ys )
        phis = increment * indices
        
        return flex.vec3_double(
            zip( flex.cos( phis ) * rs, ys, flex.sin( phis ) * rs )
            )
        

class RGIndexer(object):
    """
    Quick indexing based on ( resseq, icode )
    """
    
    def __init__(self, rgs):
        
        self.data = dict(
            [ ( ( rg.resseq_as_int(), rg.icode ), ( i, rg ) )
                for ( i, rg ) in enumerate( rgs ) ]
            )
    
    
    def __getitem__(self, key):
        
        return self.data[ key ]
    
    
    def __contains__(self, key):
        
        return key in self.data
    

class DSSPInterface(object):
    """
    Runs DSSP to determine secondary structure
    """
    
    def __init__(self, chain):
    
        import iotbx.pdb
        root = iotbx.pdb.hierarchy.root()
        model = iotbx.pdb.hierarchy.model()
        root.append_model( model )
        model.append_chain( chain.detached_copy() )
            
        import mmtbx.secondary_structure
        ( records, errors ) = mmtbx.secondary_structure.run_ksdssp_direct(
            root.as_pdb_string()
            )
        
        if errors and not records:
            raise ValueError, errors
        
        self.elements = []
        ann = iotbx.pdb.secondary_structure.process_records( records = records )
        
        if ann is not None:
            rgs = chain.residue_groups()
            indexer = RGIndexer( rgs = rgs )
            
            for helix in ann.helices:
                self.register_secondary_structure_element(
                    sse = helix,
                    rgs = rgs,
                    indexer = indexer,
                    ss_type = "helix"
                    )
                
            for sheet in ann.sheets:
                for strand in sheet.strands:
                    self.register_secondary_structure_element(
                        sse = strand,
                        rgs = rgs,
                        indexer = indexer,
                        ss_type = "sheet"
                        )
    
    
    def register_secondary_structure_element(self, sse, rgs, indexer, ss_type):
        
        si = ( sse.start_resseq, sse.start_icode )
        assert si in indexer
        ( start, rg ) = indexer[ si ]
        assert rgs[ start ] == rg
        
        ei = ( sse.end_resseq, sse.end_icode )
        assert ei in indexer
        ( end, rg ) = indexer[ ei ]
        assert rgs[ end ] == rg
        
        assert start <= end
        self.elements.append( rgs[ start : end + 1 ] )


class AtomDescription(object):
    """
    Serializable atom
    """
    
    def __init__(self, b, element, hetero, name, occ, segid, xyz):
        
        self.b = b
        self.element = element
        self.hetero = hetero
        self.name = name
        self.occ = occ
        self.segid = segid
        self.xyz = xyz
    
    
    def identifier(self):
        
        return "atom %s" % self.name
    
    
    def as_iotbx_object(self):
        
        import iotbx.pdb
        
        new_atom = iotbx.pdb.hierarchy.atom()
        new_atom.set_b( self.b )
        new_atom.set_element( self.element )
        new_atom.set_hetero( self.hetero )
        new_atom.set_name( self.name )
        new_atom.set_occ( self.occ )
        new_atom.set_segid( self.segid )
        new_atom.set_xyz( self.xyz )
        return new_atom
    
    
    def parseable_form(self):
        
        return ( self.b, self.element, self.hetero, self.name, self.occ, self.segid, self.xyz )
    
    
    def __hash__(self):
        
        return hash(
            ( self.element, self.hetero, self.name, self.occ, self.segid )
            )
    
    
    def __eq__(self, other):
        
        if not isinstance( other, self.__class__ ):
            return False
        
        return (
            ( abs( self.b - other.b ) < 0.01 )
            and ( self.element == other.element )
            and ( self.hetero == other.hetero )
            and ( self.name == other.name )
            and ( self.occ == other.occ )
            and ( self.segid == other.segid )
            and ( self.distance( point1 = self.xyz, point2 = other.xyz ) < 0.01 )
            )
        
        
    def __ne__(self, other):
        
        return not( self == other )
    
    
    @staticmethod
    def distance(point1, point2):
        
        ( x1, y1, z1 ) = point1
        ( x2, y2, z2 ) = point2
        
        return math.sqrt( ( x1 - x2 ) ** 2 + ( y1 - y2 ) ** 2 + ( z1 - z2 ) ** 2 ) 
    

class ProteinSidechainCBCompletion(object):
    """
    Calculates C-beta positions uniformly for each residue
    """
    
    C = " C  "
    CA = " CA "
    N = " N  "
    CB = " CB "
    NEEDED = ( N, CA, C )
    NAME = "cbeta"
    
    def __init__(self, atom_named, resname, max_level):
        
        if ( resname == "GLY" or max_level < 1 or self.CB in atom_named
            or any( a not in atom_named for a in self.NEEDED ) ):
            self.generated = []
            
        else:
            ca = atom_named[ self.CA ]
            
            data = AtomDescription(
                b = ca.b,
                element = " C",
                hetero = ca.hetero,
                name = self.CB,
                occ = ca.occ,
                segid = ca.segid,
                xyz = self.cb_coordinates_from(
                    ca = ca.xyz,
                    n = atom_named[ self.N ].xyz,
                    c = atom_named[ self.C ].xyz
                    ), 
                )
            self.generated = [ data ]
    
    
    @staticmethod
    def cb_coordinates_from(ca, n, c):
        
        # Residue-fixed orthonormal coordinate system (relative to CA):
        #   X along normalised N - normalised C
        #  -Y along normalised N + normalised C
        #   Z to complete a right-handed coordinate system
        n_rel = flex.vec3_double( [ n ] ) - ca
        c_rel = flex.vec3_double( [ c ] ) - ca
        x = ( 1 / n_rel.norm() ) * n_rel - ( 1 / c_rel.norm() ) * c_rel
        y = -1 * ( ( 1 / n_rel.norm() ) * n_rel + ( 1 / c_rel.norm() ) * c_rel )
        z = flex.vec3_double( [
            (
                x[0][1]*y[0][2] - x[0][2]*y[0][1],
                x[0][2]*y[0][0] - x[0][0]*y[0][2],
                x[0][0]*y[0][1] - x[0][1]*y[0][0]
                )
            ] )
        
        cb_rel = (
              0.00199448 / x.norm() * x
            + 0.92989418 / y.norm() * y
            - 1.20217920 / z.norm() * z
            )
        
        return ( cb_rel + ca )[0]
        

class SequenceDB(object):
    """
    Quick retrieval functions for matching sequences
    """
    
    def __init__(self):
        
        self.store = []
        
        
    def add(self, sequence, data):
        
        self.store.append( ( sequence, data ) )
        
        
    def matching(self, sequence):
        
        matching = []
        seq_length = len( sequence )
    
        for ( s, d ) in self.store:
            lcs = ncs.LCSTree()
            lcs.insert_word( sequence )
            lcs.insert_word( s )
            ss = lcs.get_longest_common_substring()
            matching.append( len( ss ) )
            
        return matching
    
    
    def count(self):
        
        return len( self.store )


def sequence_weight(sequence, mmt):
    """
    Calculates weight from sequence
    """
        
    if not mmt.molecular_weights or len( mmt.one_letter_codes ) != len( mmt.molecular_weights ):
        raise RuntimeError, "Incomplete data for type '%s'" % mmt.name
    
    average = sum( mmt.molecular_weights ) / len( mmt.molecular_weights )
    weight_for = dict( zip( mmt.one_letter_codes, mmt.molecular_weights ) )
    return ( sum(
        [ weight_for.get( rescode, average ) for rescode in sequence ]
        )
        - ( len( sequence ) - 1 ) * mmt.condensation_weight_loss )
    
    
def structure_weight(root, mmt):
    
    sequence = one_letter_sequence(
        rgs = root.residue_groups(),
        one_letter_for = dict( zip( mmt.three_letter_codes, mmt.one_letter_codes ) ),
        unknown = mmt.unknown_one_letter
        )
    
    return sequence_weight( sequence = sequence, mmt = mmt )