import math
from structure import *
class Atom:

    def dist_atoms(self, atom2):
        return(math.sqrt((self.coord_x-atom2.coord_x)**2 +
                         (self.coord_y-atom2.coord_y)**2 +
                         (self.coord_z-atom2.coord_z)**2))


    def __init__(self, atom_id, atom_name, res_name, chain_id,
                 res_seq_nb, insertion_code, coordinates):
        self.atom_id = atom_id
        self.atom_name = atom_name
        self.res_name = res_name
        self.chain_id = chain_id
        self.res_seq_nb = res_seq_nb
        self.insertion_code = insertion_code
        self.coord_x = coordinates[0]
        self.coord_y = coordinates[1]
        self.coord_z = coordinates[2]
        self.coords = coordinates

class Residue:
    def __init__(self, atoms_list, indice):
        self.atoms = {}
        for atom in atoms_list:
            self.atoms[atom.atom_name] = atom
            self.resid = atom.res_seq_nb
            self.res_name = atom.res_name
            self.chain_id = atom.chain_id
            self.insertion_code = atom.insertion_code
            self.indice = indice
            
            
    def h_bond(self, res2):
        if("H" not in res2.atoms.keys()):
            return(False)
        # dimensionnal factor, in kcal/mole
        f = 332
        # partial charges
        q1 = 0.42
        q2 = 0.20
        # distance between O-N atoms, in angströms
        r_ON = self.atoms["O"].dist_atoms(res2.atoms["N"])
        # distance between C-H atoms, in angströms
        r_CH = self.atoms["C"].dist_atoms(res2.atoms["H"])
        # distance between O-H atoms, in angströms
        r_OH = self.atoms["O"].dist_atoms(res2.atoms["H"])
        # distance between C-N atoms, in angströms
        r_CN = self.atoms["C"].dist_atoms(res2.atoms["N"])
        # electrostatic interaction energy, in kcal/mole
        E = q1*q2*(1/r_ON + 1/r_CH - 1/r_OH - 1/r_CN)*f
        return(E)

    def get_turns(self, residues):
        """
        Get all the turns from a specific residue.
        """
        turns = {}
        i = residues.index(self)
        k = 0
        for j in range(3,6):
                if(i+j<len(residues)):
                    if(self.h_bond(residues[i+j])<-0.5):
                        k = j
        if k != 0:    
            #print(k,"TURN", residues[i].resid, residues[i+k].resid)
            return Turn(k,residues[i].resid)

        return False
    
    def get_bridges(self, residues):
        bridges = {}
        bridge = {}
        strongest_bridge = {}
        i = residues.index(self)
        if(i >= 1 and i < len(residues)-4):
            E_min = 0
            for j in range(i+2,len(residues)-1):
                # select triplet with the minimal energy 

                if(residues[i-1].h_bond(residues[j])<-0.5
                   and residues[j].h_bond(residues[i+1])<-0.5):
                    bridge = {'res1':residues[i-1].h_bond(residues[j]),
                              'res2':residues[j].h_bond(residues[i+1]),
                              'ipos':residues[i].resid,
                              'jpos':residues[j].resid,
                              'btype':"para"}

                if(residues[j-1].h_bond(residues[i])<-0.5
                   and residues[i].h_bond(residues[j+1])<-0.5):
                    bridge = {'res1':residues[j-1].h_bond(residues[i]),
                              'res2':residues[i].h_bond(residues[j+1]),
                              'ipos':residues[i].resid,
                              'jpos':residues[j].resid,
                              'btype':"para"}

                if(residues[i].h_bond(residues[j])<-0.5
                   and residues[j].h_bond(residues[i])<-0.5):
                    bridge = {'res1':residues[i].h_bond(residues[j]),
                              'res2':residues[j].h_bond(residues[i]),
                              'ipos':residues[i].resid,
                              'jpos':residues[j].resid,
                              'btype':"anti"}

                if(residues[i-1].h_bond(residues[j+1])<-0.5
                   and residues[j-1].h_bond(residues[i+1])<-0.5):
                    bridge = {'res1':residues[i-1].h_bond(residues[j+1]),
                              'res2':residues[j-1].h_bond(residues[i+1]),
                              'ipos':residues[i].resid,
                              'jpos':residues[j].resid,
                              'btype':"anti"}

                if(bridge):
                    if(bridge['res1']+bridge['res2']<E_min):
                        E_min = bridge['res1']+bridge['res2']
                        strongest_bridge = bridge
                        bridge = {}
                        coord_bridge = [i,j]
            # finally add the strongest bridge at i and j pos
            if(strongest_bridge):
                bridges[coord_bridge[0]] = (Bridge(strongest_bridge['btype'],
                                                   strongest_bridge['ipos'],
                                                   strongest_bridge['jpos']))
                
                bridges[coord_bridge[1]] = (Bridge(strongest_bridge['btype'],
                                                   strongest_bridge['jpos'],
                                                   strongest_bridge['ipos']))
        if(len(bridges)>0):
            return(bridges[coord_bridge[0]])
        else:
            return(False)

    def get_helix(self, residues):
        """
        Return if there is an helix at a given residue,
        as well as its type.
        """
        i = residues.index(self)
        # if there are no turns or it is the first residue, skip
        if i == 0:
            return False
        
        if(self.get_turns(residues) and residues[i-1].get_turns(residues)):
            print(self.get_turns(residues).turn_type,"- HELIX at", residues[i].resid)
            return(self.get_turns(residues).turn_type, residues[i].resid)
        return(False)

    def get_ladder(self, residues):
        #ladders = {}
        i = residues.index(self)
        if i != 0:
            if self.get_bridges(residues):
                if (residues[i-1].get_bridges(residues)):
                    local_bridge = self.get_bridges(residues)
                    consec_bridge = residues[i-1].get_bridges(residues)
                    if local_bridge.bridge_type == consec_bridge.bridge_type:
                        #print("ladder", consec_bridge.res_num, local_bridge.res_num)
                        ladder = {'start':consec_bridge.res_num,
                                  'end':local_bridge.res_num,
                                  'bridges':[consec_bridge, local_bridge]}
                        return ladder
        return False   

def get_bridges(residues):
    bridges = {}
    bridge = {}
    strongest_bridge = {}
    for i in range(1,len(residues)-4):
        E_min = 0
        for j in range(i+2,len(residues)-1):
            # select triplet with the minimal energy 

            if(residues[i-1].h_bond(residues[j])<-0.5
               and residues[j].h_bond(residues[i+1])<-0.5):
                bridge = {'res1':residues[i-1].h_bond(residues[j]),
                          'res2':residues[j].h_bond(residues[i+1]),
                          'ipos':residues[i].resid,
                          'jpos':residues[j].resid,
                          'i':residues[i].indice,
                          'j':residues[j].indice,
                          'btype':"para"}
                   
            if(residues[j-1].h_bond(residues[i])<-0.5
               and residues[i].h_bond(residues[j+1])<-0.5):
                bridge = {'res1':residues[j-1].h_bond(residues[i]),
                          'res2':residues[i].h_bond(residues[j+1]),
                          'ipos':residues[i].resid,
                          'jpos':residues[j].resid,
                          'i':residues[i].indice,
                          'j':residues[j].indice,
                          'btype':"para"}

            if(residues[i].h_bond(residues[j])<-0.5
               and residues[j].h_bond(residues[i])<-0.5):
                bridge = {'res1':residues[i].h_bond(residues[j]),
                          'res2':residues[j].h_bond(residues[i]),
                          'ipos':residues[i].resid,
                          'jpos':residues[j].resid,
                          'i':residues[i].indice,
                          'j':residues[j].indice,
                          'btype':"anti"}
                    
            if(residues[i-1].h_bond(residues[j+1])<-0.5
               and residues[j-1].h_bond(residues[i+1])<-0.5):
                bridge = {'res1':residues[i-1].h_bond(residues[j+1]),
                          'res2':residues[j-1].h_bond(residues[i+1]),
                          'ipos':residues[i].resid,
                          'jpos':residues[j].resid,
                          'i':residues[i].indice,
                          'j':residues[j].indice,
                          'btype':"anti"}

            if(bridge):
                if(bridge['res1']+bridge['res2']<E_min):
                    E_min = bridge['res1']+bridge['res2']
                    strongest_bridge = bridge
                    coord_bridge = [i,j]
                    bridge = {}
        # finally add the strongest bridge at i and j pos
        if(strongest_bridge):
            bridges[strongest_bridge['i']] = (Bridge(strongest_bridge['btype'],
                                                     strongest_bridge['ipos'],
                                                     strongest_bridge['jpos'],
                                                     [strongest_bridge['i'],
                                                      strongest_bridge['j']]))
            bridges[strongest_bridge['j']] = (Bridge(strongest_bridge['btype'],
                                               strongest_bridge['jpos'],
                                               strongest_bridge['ipos'],
                                               [strongest_bridge['i'],
                                               strongest_bridge['j']]))
    if(len(bridges)>0):
        return(bridges)
    else:
        return(False)

def get_ladders(bridges, residues):
    ladders = {}
    i = 1
    while i < len(residues):
        k = 1
        if i in bridges.keys():
            temp_bridges = [bridges[i]]
            while ((i+k in bridges.keys()) and
                   (bridges[i].bridge_type == bridges[i+k].bridge_type)):
                temp_bridges.append(bridges[i+k])
                k+=1
        if k>1:
            #print("ladder", bridges[i].res_num, bridges[i+k-1].res_num)
            ladders[i] = {'start':bridges[i].res_num,
                          'end':bridges[i+k-1].res_num,
                          'bridges':temp_bridges,
                          'i':i,
                          'j':i+k-1}
            i+=k-1
        else:
            i+=1
    return ladders

def connected_ladders(ladd_1, ladd_2):
    links = []
    for bridge in ladd_1['bridges']:
        if bridge.res_partner in res_list(ladd_2):
            return([ladd_1['i'], ladd_1['j'], bridge.i, bridge.j,
                    ladd_2['i'], ladd_2['j']])

    return False

def get_sheets(ladders):
    """
    Bridges between ladders.
    Check if 1 bridge between one ladder and one or more other ladders.
    Iterate over all residues of one ladder and check if bridge with other residues
    of the other ladders.
    """
    ladds = [ elem for elem in ladders.values() ]
    sheets = {}
    i = 0
    while i < len(ladds):
        j = i+1
        ladd = ladds[i]
        ladd1 = ladds[i]
        ladd2 = ladds[j]

        while j < len(ladds):
            print(connected_ladders(ladd1, ladd2))
            if connected_ladders(ladd1, ladd2)!=False:
                ladd1 = ladd2
                #print(connected_ladders(ladd1, ladd2)[4])
                ladd2 = ladders[connected_ladders(ladd1, ladd2)[4]]
            j+=1
        print(ladd['i'], ladd2['i'])
        i+=1
    # for ladder in ladders:
    #     for ladd2 in ladders:
    #        if connected_ladders(ladders[ladder], ladders[ladd2]):
    #            pass
               #print("ladder",ladders[ladder]['i'], ladders[ladder]['j'],"bridge",bridge.i, bridge.j,
               #      "ladder 2",ladders[ladd2]['i'], ladders[ladd2]['j'])

def res_list(ladder):
    # TODO : method in ladder class
    l=[]
    for i in range(ladder['start'], ladder['end']):
        l.append(i)
    return(l)
                   
def get_bends(residues):
    bends = {}
    for i in range(2,len(residues)-2):
        angle = math.degrees(vector_angles(vectors_substr(position_vector(residues[i].atoms["CA"].coords),
                                                          position_vector(residues[i-2].atoms["CA"].coords)),
                                           vectors_substr(position_vector(residues[i+2].atoms["CA"].coords),
                                                          position_vector(residues[i].atoms["CA"].coords))))
        if(angle>70):
            print("angle", residues[i].resid, angle)
            bends[residues[i].resid] = angle
    return(bends)

def get_chirality(residues):

    for i in range(1,len(residues)-2):
        chirality = {}
        angle = calc_dihedral(residues[i-1].atoms["CA"].coords,
                              residues[i].atoms["CA"].coords,
                              residues[i+1].atoms["CA"].coords,
                              residues[i+2].atoms["CA"].coords)
        
        if(angle>0 and angle<=180):
            sign="+"
            print("chirality", residues[i].resid, "+", angle)
            
        if(angle<=0 and angle > -180):
            sign="-"
            print("chirality", residues[i].resid, "-", angle)
            
        chirality[residues[i].resid] = sign
        
    return chirality


def get_phi_psi(residues):
    for i in range(len(residues)):
        if(i==0):
            phi = 360.0
        else:
            phi = calc_dihedral(residues[i-1].atoms["C"].coords,
                                residues[i].atoms["N"].coords,
                                residues[i].atoms["CA"].coords,
                                residues[i].atoms["C"].coords)
        if(i==len(residues)-1):
            psi = 360.0
        else:
            psi = calc_dihedral(residues[i].atoms["N"].coords,
                                residues[i].atoms["CA"].coords,
                                residues[i].atoms["C"].coords,
                                residues[i+1].atoms["N"].coords)

        print("ANGLES", i, phi, psi)

def get_TCO(res1, res2):
    CO_res1 = vector_from_pos(res1.atoms["C"].coords,
                              res1.atoms["O"].coords)
    CO_res2 = vector_from_pos(res2.atoms["C"].coords,
                              res2.atoms["O"].coords)
    angle = vector_angles(CO_res1, CO_res2)
    return(math.cos(angle))
