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[residues[i].resid] = 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[i-1].h_bond(residues[j])+
                #    residues[j].h_bond(residues[i+1]))<E_min:
                #    E_min = residues[i-1].h_bond(residues[j])
                #    +residues[j].h_bond(residues[i+1])
                #    bridge_type = "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[j-1].h_bond(residues[i])+
                #    residues[i].h_bond(residues[j+1]))<E_min:
                #    E_min = residues[j-1].h_bond(residues[i])
                #    +residues[i].h_bond(residues[j+1])
                #    bridge_type = "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].h_bond(residues[j])+
                #    residues[j].h_bond(residues[i]))<E_min:
                #    E_min = residues[i].h_bond(residues[j])
                #    +residues[j].h_bond(residues[i])
                #    bridge_type = "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(residues[i-1].h_bond(residues[j+1])+
                 #   residues[j-1].h_bond(residues[i+1]))<E_min:
                 #   E_min = residues[i-1].h_bond(residues[j+1])
                 #   +residues[j-1].h_bond(residues[i+1])
                 #   bridge_type = "anti"
            if(bridge):
                if(bridge['res1']+bridge['res2']<E_min):
                    E_min = bridge['res1']+bridge['res2']
                    strongest_bridge = bridge
                    bridge = {}
        # finally add the strongest bridge at i and j pos
        if(strongest_bridge):
            bridges[strongest_bridge['ipos']] = (Bridge(strongest_bridge['btype'],
                                                        strongest_bridge['ipos'],
                                                        strongest_bridge['jpos']))
            bridges[strongest_bridge['jpos']] = (Bridge(strongest_bridge['btype'],
                                                        strongest_bridge['jpos'],
                                                        strongest_bridge['ipos']))
    if(len(bridges)>0):
        return(bridges)
    else:
        return(False)

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])
            if(k>2):
                print(k,"- 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):
    ladders = {}
    i = 1
    while i < len(bridges):
        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[bridges[i].res_num] = k-1
            print("ladder", bridges[i].res_num, bridges[i+k-1].res_num)
            ladders[bridges[i].res_num] = {'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.
    """

    for ladder in ladders:
        for bridge in ladders[ladder]['bridges']:
            for ladd2 in ladders:
                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'])
                

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 = 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)
    deg_angle = rad_angle*(180/math.pi)
    return deg_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