# Simple topology representation for sidechains
ALA = (
    ( " CB ", () ),
    )
CYS = (
    ( " CB ", (
        ( " SG ", () ),
        ), ),
    )
ASP = (
    ( " CB ", (
        ( " CG ", (
            ( " OD1", () ),
            ( " OD2", () ),
            ), ),
        ), ),
    )
GLU = (
    ( " CB ", (
        ( " CG ", (
            ( " CD ", (
                ( " OE1", () ),
                ( " OE2", () ),
                ), ),
            ), ),
        ), ),
    )
PHE = (
    ( " CB ", (
        ( " CG ", (
            ( " CD1", (
                ( " CE1", (
                    ( " CZ ", () ), 
                    ), ),
                ), ),
            ( " CD2", (
                ( " CE2", () ),
                ), ),
            ), ),
        ), ),
    )
GLY = ()
HIS = (
    ( " CB ", (
        ( " CG ", (
            ( " ND1", (
                ( " CE1", () ),
                ), ),
            ( " CD2", (
                ( " NE2", () ),
                ), ),
            ), ),
        ), ),
    )
ILE = (
    ( " CB ", (
        ( " CG1", (
            ( " CD1", () ),
            ), ),
        ( " CG2", () ),
        ), ),
    )
LYS = (
    ( " CB ", (
        ( " CG ", (
            ( " CD ", (
                ( " CE ", (
                    ( " NZ ", () ),
                    ), ),
                ), ),
            ), ),
        ), ),
    )
LEU = (
    ( " CB ", (
        ( " CG ", (
            ( " CD1", () ),
            ( " CD2", () ),
            ), ),
        ), ),
    )
MET = (
    ( " CB ", (
        ( " CG ", (
            ( " SD ", (
                ( " CE ", () ),
                ), ),
            ), ),
        ), ),
    )
ASN = (
    ( " CB ", (
        ( " CG ", (
            ( " OD1", () ),
            ( " ND2", () ),
            ), ),
        ), ),
    )
PRO = (
    ( " CB ", (
        ( " CG ", (
            ( " CD ", () ),
            ), ),
        ), ),
    )
GLN = (
    ( " CB ", (
        ( " CG ", (
            ( " CD ", (
                ( " OE1", () ),
                ( " NE2", () ),
                ), ),
            ), ),
        ), ),
    )
ARG = (
    ( " CB ", (
        ( " CG ", (
            ( " CD ", (
                ( " NE ", (
                    ( " CZ ", (
                        ( " NH1", () ),
                        ( " NH2", () ),
                        ), ),
                    ), ),
                ), ),
            ), ),
        ), ),
    )
SER = (
    ( " CB ", (
        ( " OG ", () ),
        ), ),
    )
THR = (
    ( " CB ", (
        ( " OG1", () ),
        ( " CG2", () ),
        ), ),
    )
VAL = (
    ( " CB ", (
        ( " CG1", () ),
        ( " CG2", () ),
        ), ),
    )
TRP = (
    ( " CB ", (
        ( " CG ", (
            ( " CD2", (
                ( " CE2", (
                    ( " CZ2", (
                        ( " CH2", () ),
                        ), ),
                    ), ),
                ( " CE3", (
                    ( " CZ3", () ),
                    ), ),
                ), ),
            ( " CD1", (
                ( " NE1", () ),
                ), ),
            ), ),
        ), ),
    )
TYR = (
    ( " CB ", (
        ( " CG ", (
            ( " CD1", (
                ( " CE1", (
                    ( " CZ ", ( 
                        ( " OH ", () ),
                        ), ),
                    ), ),
                ), ),
            ( " CD2", (
                ( " CE2", () ),
                ), ),
            ), ),
        ), ),
    )


class SidechainAtom(object):
    """
    Calculates atom names from topology
    """
    
    def __init__(self, topology):
        
        self.names = self.flatten( subtree = topology )
        
        
    @classmethod
    def flatten(cls, subtree):
        
        atoms = []
        
        for entry in subtree:
            atoms.extend(
                [ entry[0] ] 
                + cls.flatten( entry[1] ) )
            
        return atoms
        

class MainchainDistance(object):
    """
    Calculates distances from mainchain from topology
    """
    
    def __init__(self, topology, names):
        
        self.from_atom = dict(
            [
                (
                    atom,
                    len( self.get_path_from( tree = topology, item = atom ) ),
                    )
                for atom in names
                ]
            )
    
    
    @classmethod
    def get_path_from(cls, tree, item):
        
        if any( [ entry[0] == item for entry in tree ] ):
            return [ item ]
        
        for entry in tree:
            trial_path = cls.get_path_from( entry[-1], item )
            
            if trial_path:
                return [ entry[0] ] + trial_path
            
        return []
    
    
    @classmethod
    def All(cls, topology):
        
        sca = SidechainAtom( topology = topology )
        return cls( topology = topology, names = sca.names )

    
class SidechainCorrespondence(object):
    """
    Creates a mapping from old residue atoms to new residue atoms
    """
    
    def __init__(self, source, target):
        
        self.mapping = dict(
            self.create_mapping( source = source, target = target )
            )
    
    
    @classmethod
    def create_mapping(cls, source, target):
        
        mapping = []
        
        for ( s, t ) in zip( source, target ):
            mapping.append( ( s[0], t[0] ) )
            mapping.extend( cls.create_mapping( source = s[1], target = t[1] ) )
            
        return mapping