#!/usr/bin/python

"""Calculate the eigenvalues and eigeinvectors from residues to
determine the rotation of the protein.
"""

import os
import numpy as np
import math
import matplotlib.pyplot as plt


__author__ = ""
__date__ = "14/11/14"
__version__ = "1.0.0"
__copyright__ = "CC_by_SA"
__dependencies__ = "os, numpy, math, matplotlib"



def frange(start, stop, step, dec = 3):
    """Range like function accepting float step

    Args:
        start (int)
        stop (int)
        step (int)
        dec (int) number of decimal

    Return:
        A list number from start to stop with the given step
    """
    f_tab = []
    val = start

    while(val<=stop):
        f_tab.append(val)
        val = round(val + step, dec)

    return f_tab



def parse_pdb(filepath):
    """Parse pdb file and return numpy matrix with atoms positions

    Args:
        filepath (str) path to the .pdb file

    Retrun:
        A numpy matrix containing atoms positions (x, y, z as column).

    """
    i = 0;

    with open(filepath, 'r') as pdbfile:

        for line in pdbfile:
            if line[:4] == "ATOM":

                x_pos = float(line[30:38])
                y_pos = float(line[38:46])
                z_pos = float(line[46:54])

                if (i==0):
                    positions = np.array([x_pos,y_pos,z_pos])
                    i += 1;

                else:
                    positions = np.vstack([positions, [x_pos, y_pos, z_pos]]);

    return positions;



def eigen_vectors(pdbfile):
    """Return the eigenvectors from a pdb file

    Args:
        pdbfile (str) a path to the pdbfile to parse

    Return:
        A numpy vector containing the eigenvectors (length from origin = 1).
    """

    # Get the matrix of x,y,z position
    pos_mtx = parse_pdb(pdbfile);

    pos_mtx[:,0] = pos_mtx[:,0]- np.mean(pos_mtx[:,0])
    pos_mtx[:,1] = pos_mtx[:,1]- np.mean(pos_mtx[:,1])
    pos_mtx[:,2] = pos_mtx[:,2]- np.mean(pos_mtx[:,2])
    # Calculate the inertia matrix multiplying the transpose matrix to itself
    square_mtx = np.transpose(pos_mtx).dot(pos_mtx)

    # eigenvectors
    eig_vec = np.linalg.eig(square_mtx)[1][:,0]

    return eig_vec;



def zaxis_angles(axis):
    """Determine the angle between z-axis and the axis specified

    Args:
        axis (numpy vector) the normalized (length=1) vectors to calculate the
        angle

    Return:
        A tuple containing angles values in radian unit
    """

    angle = []
    # z axis vector
    z_vec = np.array([0,0,1])

    # calculate angle in radian
    angle.append(math.acos(z_vec.dot(axis))*57.3)

    return angle;



def z_orientations(filepath):
    """Return a list of rotation from a directory containing frame's
    trajectory as .pdb files

    Args:
        filepath (str) path to the directory containing trajectory .pdb files

    Return:
        list of angles from the z_axis for each frame
    """

    cor_idx = cor_eigvec(filepath + "/traj0.pdb")

    nb_fichier = len(os.listdir(filepath))

    for i in xrange(nb_fichier):
        pdb = filepath + "/traj%i.pdb"%i

        eig_vec = eigen_vectors(pdb)
        atoms = parse_pdb(pdb)

        cor_vec = vecteur_unitaire(atoms, cor_idx[0], cor_idx[1])

        same_direction(eig_vec, cor_vec)




def vecteur_unitaire(atoms, i, j):
    """
    """

    vect_cor = atoms[j,:] - atoms[i,:]

    distance = math.sqrt((atoms[j,0]-atoms[i,0])**2+
                         (atoms[j,1]-atoms[i,1])**2+
                         (atoms[j,2] - atoms[i,2])**2)

    vect_cor = vect_cor / distance;
    return vect_cor;


def same_direction(eig_vec, vec_cor):

    scal = vec_cor.dot(eig_vec)

    if scal > 0:
        return eig_vec

    else:
        return -eig_vec


def cor_eigvec(filepath, min_dis=17, correl = 0.9):
    """
    """

    atoms = parse_pdb(filepath)
    eig_vec = eigen_vectors(filepath)
    is_ok = False


    for i in xrange(len(atoms)):
        for j in xrange(len(atoms)):

            if (i==j):
                continue;

            vect_cor = atoms[j,:] - atoms[i,:]

            distance = math.sqrt((atoms[j,0]-atoms[i,0])**2+
                            (atoms[j,1]-atoms[i,1])**2+
                            (atoms[j,2] - atoms[i,2])**2)


            if (distance <= min_dis):
                continue;

            vect_cor = vect_cor/distance
            scal = vect_cor.dot(eig_vec)

            if (scal >= correl):
                is_ok = True
                break

        if is_ok == True:
            break

    if i==j and i==len(atoms[:,0]):
        print "Attention, pas de vecteur de correction trouve pour ces valeurs de distance et de correlation"
        return (-1, -1);

    return (i, j);


if __name__ == "__main__":

    ori = z_orientations("pdbs/");

    time = frange(0, 20, 0.1) # 101 frame avec un pas de 0.2 nm.

    plt.plot(time, ori)
    plt.show()