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):
        self.bridge_type = bridge_type
        self.res_num = res_num
        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):
    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.append(Bridge(strongest_bridge['btype'],
                                  strongest_bridge['ipos']))
            bridges.append(Bridge(strongest_bridge['btype'],
                                  strongest_bridge['jpos']))
    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_ladder(bridges):
    ladders = {}
    for i in range(len(bridges)):
        k = 1
        while bridges[i] == bridges[i+k]:
            k+=1
        if k>1:
            ladders[bridges[i].res_num] = k-1

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("angle", residues[i].resid, "+", angle)
            
        if(angle<=0 and angle > -180):
            sign="-"
            print("angle", residues[i].resid, "-", angle)
            
        chirality[residues[i].resid] = sign
        
    return chirality