#!/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).

    """

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

                try:
                    positions = np.vstack((positions, [x_pos, y_pos, z_pos]));
                except:
                    positions = np.array([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]
    print eig_vec;
    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)
    #angles.append(math.acos(z_vec.dot(axis[:,1]))*57.3)
    #angles.append(math.acos(z_vec.dot(axis[:,2]))*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
    """

    files = os.listdir(filepath)

    orientation = []
    for pdb in files:
        eig_vec = eigen_vectors(filepath + pdb)
        angle = zaxis_angles(eig_vec)

        orientation.append(angle)
        #orientations[1].append(angles[1])
        #orientations[2].append(angles[2])

    return orientation;



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