import re

class Distances(object):
    '''
    Store and access distance data for residues from a protein structure.
    '''
    def __init__(self, dist_data):
        '''
        Creates a new instance from distance data.
        
        Distance data should be a list of tuples of five elements: (chain1, pos1, chain2, pos2, distance).
        :param dist_data: a list of (chain1, pos1, chain2, pos2, distance)
        '''
        dis = {}
        for c1, p1, c2, p2, d in dist_data:
            if not (c1, p1) in dis:
                dis[(c1, p1)] = {}
            dis[(c1, p1)][(c2,p2)] = d 
        self._distances = dis
#         self._distances = {(c1, p1): {(c2,p2): d
#             for c1p, p1p, c2, p2, d in dist_data if c1p == c1 and p1p == p1}
#             for c1, p1, _, _, _ in dist_data}
    
    def of(self, chain_a, pos_a, chain_b, pos_b):
        '''
        Retrieves distance for a residue pair.
        
        If the pair is not found, None is returned.
        :param chain_a: A string specifying the first residue chain.
        :param pos_a: An integer specifying the first residue position.
        :param chain_b: A string specifying the second residue chain.
        :param pos_b: An integer specifying the second residue position.
        '''
        pair1 = ((chain_a, pos_a))
        pair2 = ((chain_b, pos_b))
        if pair1 == pair2: # Special case for distance with the same residue.
            return 0
        d = self._distances.get( pair1 , {}).get(pair2)
        if not d:
            d = self._distances.get( pair2 , {}).get(pair1)
        return d
    
    def remap_positions(self, mapping):
        '''
        Remap index positions.
        
        If a positions could not be mapped it is excluded from the results.
        :param mapping: a dict that maps old positions to new positions.
        '''
        def remap(dic):
            return {(chain, mapping[chain][pos]):value 
                    for (chain, pos), value in dic.items()
                    if pos in mapping.get(chain, {})}
        
        self._distances = remap({(c1,p1):remap(r2) 
                                 for (c1,p1), r2 in self._distances.items()})
        
    def is_contact(self, chain_a, pos_a, chain_b, pos_b, distance_cutoff = 6.05):
        '''
        Returns True if a given pair's distance is lower or equal than a given
        distance cutoff. 
        :param chain_a: A string specifying the first residue chain.
        :param pos_a: An integer specifying the first residue position.
        :param chain_b: A string specifying the second residue chain.
        :param pos_b: An integer specifying the second residue position.
        :param distance_cutoff: a float with the distance cutoff (defaults to 6.05 angstroms)
        '''
        return self.of(chain_a, pos_a, chain_b, pos_b) <= distance_cutoff
        

def from_mitos(dist_file):
    '''
    Loads data of residue distances from a file generated by MIToS.
    
    Input data should look like:
    
    <pre>
    # model_i,chain_i,group_i,pdbe_i,number_i,name_i,model_j,chain_j,group_j,pdbe_j,number_j,name_j,distance
    1,A,ATOM,,55,LEU,1,A,ATOM,,56,LEU,1.3247309160731473
    </pre>

    :param dist_file: A string to a text file with the distance data.
    '''
    # model_i,chain_i,group_i,pdbe_i,number_i,name_i,model_j,chain_j,group_j,pdbe_j,number_j,name_j,distance
    # 1,A,ATOM,,55,LEU,1,A,ATOM,,56,LEU,1.3247309160731473
    #                       1    ,A  ,ATOM,  ,55   ,LEU  ,1   ,A  ,ATOM,  ,56   ,LEU ,1.3247309160731473
    d_pattern = re.compile("(\d+),(.),(.+),.*,(\d+),(.+),(\d+),(.),(.+),.*,(\d+),(.+),(.+)$")
    res = []
    with open(dist_file) as fh:
        for line in fh:
            line = line.strip()
            if not line.startswith("#"):
                m = re.match(d_pattern, line)
                try:
                    res.append((
                        m.group(2),      # Chain 1 
                        int(m.group(4)), # Pos res 1 
                        m.group(7),      # Chain 2
                        int(m.group(9)), # Pos res 2
                        float(m.group(11))))  # distance
                except:
                    pass
        return res
