import math
import numpy as np

class Structure:
    def __init__(self, res):
        self.res = res

class Turn(Structure):

    def __init__(self, turn_type, res):
        self.turn_type = turn_type
        Structure.res = res
                        
class Bridge(Structure):

    def __init__(self, bridge_type, res_num, res_partner):
        self.bridge_type = bridge_type
        self.res_num = res_num
        self.res_partner = res_partner
        Structure.res = res_num

class Helix(Structure):

    def __init__(self, residues, res_num):
        self.residues = residues
        self.res_num = res_num
        Structure.res = res_num
        
def get_turns(residues):
    # TODO : prevent redondency of overlapping turns
    turns = {}
    for i,res in enumerate(residues):
        for j in range(3,6):
            if(i+j<len(residues)):
                if(res.h_bond(residues[i+j])<-0.5):
                    print(j,"TURN", residues[i].resid, residues[i+j].resid)
                    turns[i] = Turn(j,residues[i].resid)
    return(turns)

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,
                          '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)
    else:
        return(False)

# def get_bonds(residues):
#     for i,res1 in enumerate(residues):
#         E_min = 0
#         strongest_bridge = []
#         for j in range(-5,6):
#             if(i+j < len(residues)):
#                 res2 = residues[i+j]
#                 if(res1.h_bond(res2)<-0.5) and (res1.h_bond(res2)<E_min):
#                     E_min = res1.h_bond(res2)
#                     strongest_bridge = [res1, res2]
#         if(len(strongest_bridge)>0):
#             diff = strongest_bridge[0].resid - strongest_bridge[1].resid
#             if( abs (diff) > 1 and abs(diff)<=5):
#                 print(diff)

def get_bonds(residues):
    bonds = {}
    k = 0
    for i,res in enumerate(residues):
        E_min = 0
        for j in range(-5,-2):
            if(i+j<len(residues)):
                if(residues[i+j].h_bond(res)<-0.5):
                    k+=1
                    #if(res.h_bond(residues[i+j])<E_min):
                    E_min = res.h_bond(residues[i+j])
                    res_a = res
                    res_b = residues[i+j]
                    if j not in bonds:
                        bonds[j] = []
                    bonds[j].append([res_a.resid, res_b.resid])
                    #turns[residues[i].resid] = Turn(j,residues[i].resid)
    for key in bonds.keys():
        print(key, len(bonds[key]))

    print("LH",k)
    
def get_helix(residues, turns):
    i = 1
    helix = []
    while i <= len(residues):
        if(i in turns.keys() and i-1 in turns.keys()):
            k = 0
            temp_res = []
            while(i+k in turns):
                k+=1
                temp_res.append(residues[i])
                last_res_pos = residues[i].resid+k
            if(k>=1):
                helix_size=last_res_pos - residues[i].resid
                print(turns[i].turn_type,"- HELIX at", residues[i].resid)
                helix.append(Helix(temp_res,residues[i].resid))
            i = i+k
        else:
            i+=1
    return(helix)

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:
            ladders[i] = 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+=k-1
        else:
            i+=1
    return ladders
 
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.
    """
    sheets = {}
    for ladder in ladders:
        for ladd2 in ladders:
            for bridge in ladders[ladder]['bridges']:
                if bridge.res_partner in res_list(ladders[ladd2]):
                    print("ladder",ladders[ladder]['start'], ladders[ladder]['end'],"bridge",bridge.res_num, bridge.res_partner,
                    "ladder 2",ladders[ladd2]['start'], ladders[ladd2]['end'])
            #print("stop ladder 2")
        print("stop ladder 1")            

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 vector_from_pos(a, b):
    xAB = b[0]-a[0]
    yAB = b[1]-a[1]
    zAB = b[2]-a[2]
    coord_AB = [xAB,yAB,zAB]
    return coord_AB

def vector_norm(v):
    norm = math.sqrt(v[0]**2 + v[1]**2 + v[2]**2)
    return norm

def dot_product(v1,v2):
    dot_product = v1[0]*v2[0] + v1[1]*v2[1] + v1[2]*v2[2]
    return(dot_product)

def position_vector(c):
    vector = vector_from_pos([0,0,0],c)
    return vector

def vectors_substr(v1, v2):
    return ([v1[0]-v2[0],
             v1[1]-v2[1],
             v1[2]-v2[2]])

def vector_angles(v1,v2):
    dot_prod = dot_product(v1,v2)
    norm_v1 = vector_norm(v1)
    norm_v2 = vector_norm(v2)
    term = dot_prod/(abs(norm_v1)*abs(norm_v2))
    rad_angle = math.acos(term)
    return rad_angle

def calc_dihedral(u1, u2, u3, u4):
    """ Calculate dihedral angle method. From bioPython.PDB
    (adapted to np.array)
    Calculate the dihedral angle between 4 vectors
    representing 4 connected points. The angle is in
    [-pi, pi].

    Adapted function of dihedral_angle_from_coordinates.py 
    by Eric Alcaide.
    Source : https://gist.github.com/EricAlcaide
    URL : https://gist.github.com/EricAlcaide/30d6bfdb8358d3a57d010c9a501fda56
    """
    #convert coords to numpy arrays
    u1 = np.array(u1)
    u2 = np.array(u2)
    u3 = np.array(u3)
    u4 = np.array(u4)
    
    a1 = u2 - u1
    a2 = u3 - u2
    a3 = u4 - u3

    v1 = np.cross(a1, a2)
    v1 = v1 / (v1 * v1).sum(-1)**0.5
    v2 = np.cross(a2, a3)
    v2 = v2 / (v2 * v2).sum(-1)**0.5
    porm = np.sign((v1 * a3).sum(-1))
    rad = np.arccos((v1*v2).sum(-1) / ((v1**2).sum(-1) * (v2**2).sum(-1))**0.5)
    if not porm == 0:
        rad = rad * porm

    deg_angle = rad*(180/math.pi)
    return deg_angle


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))