import numpy as np
import scipy.optimize

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt

def fitPlaneLTSQ(XYZ):
    (rows, cols) = XYZ.shape
    G = np.ones((rows, 3))
    G[:, 0] = XYZ[:, 0]  #X
    G[:, 1] = XYZ[:, 1]  #Y
    Z = XYZ[:, 2]
    (a, b, c),resid,rank,s = np.linalg.lstsq(G, Z, rcond=None)
    normal = (a, b, -1)
    nn = np.linalg.norm(normal)
    normal = normal / nn
    return (c, normal)

def angle_between2D(p1, p2):
    #arctan2 is anticlockwise
    ang1 = np.arctan2(*p1[::-1])
    ang2 = np.arctan2(*p2[::-1])
    angle=np.rad2deg((ang1 - ang2) % (2 * np.pi)) #degree
    if angle<180:
        return angle
    else:
        return 360 - angle

def angle_between3D(v1, v2):
# v1 is your firsr vector
# v2 is your second vector
    angle = np.arccos(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))) #rad
    angle=180*angle/np.pi #degrees
    if angle<180:
        return angle
    else:
        return 360 - angle

def plot_surf(data, normal):
    #Plot surface
    fig = plt.figure()
    ax = fig.gca(projection='3d')

    # plot fitted plane
    maxx = np.max(data[:,0])
    maxy = np.max(data[:,1])
    minx = np.min(data[:,0])
    miny = np.min(data[:,1])

    point = np.array([0.0, 0.0, c])
    d = -point.dot(normal)

    # plot original points
    ax.scatter(data[:, 0], data[:, 1], data[:, 2])

    # compute needed points for plane plotting
    xx, yy = np.meshgrid([minx, maxx], [miny, maxy])
    z = (-normal[0]*xx - normal[1]*yy - d)*1. / normal[2]

    # plot plane
    ax.plot_surface(xx, yy, z, alpha=0.2)

    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')
    plt.show()

def count_frames(trajfile='traj.trr'):
    from pytrr import (
        read_trr_header,
        read_trr_data,
        skip_trr_data,
        )

    cnt_fr=0
    with open(trajfile, 'rb') as inputfile:
        for i in range(1000):
            try: 
                header = read_trr_header(inputfile)
                #print('Step: {step}, time: {time}'.format(**header))
                skip_trr_data(inputfile, header)
                cnt_fr=cnt_fr+1
            except EOFError:
                pass
    return cnt_fr

def last_fr_export(trajfile='traj.trr'):
    from pytrr import (
        read_trr_header,
        read_trr_data,
        skip_trr_data,
        )
    cnt_fr=count_frames(trajfile)
    with open(trajfile, 'rb') as inputfile:
        for i in range(cnt_fr-1):
            header = read_trr_header(inputfile)
            #print('Step: {step}, time: {time}'.format(**header))
            skip_trr_data(inputfile, header)
        header = read_trr_header(inputfile)
        print('Step: {step}, time: {time}'.format(**header))
        data = read_trr_data(inputfile, header)
        print(data.keys())
        print(data['box'])
        print(data['x'][0])

#Random data
#data = np.random.randn(100, 3)/3
#data[:, 2] /=10

#Find surface of points-> for lipid plane
#c, normal = fitPlaneLTSQ(data)
#print('Normal of surface is: ',normal)

#Angle between two lines-> for finding tilt
#P1=(normal[0],normal[1], normal[2])
#P2=(0,0,-0.99998042)

#print('Line is: ',P2)
#print('Angle between vectors is: ',angle_between3D(P1,P2))

#Fit line to points->for lipid chain
#from scipy import stats
#slope, intercept, r_value, p_value, std_err = stats.linregress(data[:,0],data[:,1])
#print('y=',slope,'x+',intercept)
#print('r_value',r_value)

#Angle between two lines-> for finding tilt
#P1=(normal[0],normal[1])
#P2=(1,1)

#print('Line is: ',P2)
#print('Angle between lines is: ',angle_between2D(P1,P2))

#plot_surf(data=data, normal=normal)