import numpy as np
import scipy.optimize

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

from pytrr import (
    read_trr_header,
    read_trr_data,
    skip_trr_data,
    )

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, c):
    #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'):
    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 fr_export(trajfile='traj.trr', num_frames=1):
    cnt_fr=count_frames(trajfile)
    with open(trajfile, 'rb') as inputfile:
        for i in range(cnt_fr-num_frames):
            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])

def read_gro_types(gro):
    cnt=0
    data_num=0
    with open(gro, 'r') as F:
        for line in F:
            cnt=cnt+1
            print(cnt)
            if cnt>2:
                res_num = line[:5]  
                res_type = line[5:10]  
                atom_type = line[10:15] 
                atom_num  = line[15:20]
                rest_dt = line[20:]
                print(res_num,res_type,atom_type,atom_num,rest_dt)
            elif cnt==1:
                system=line[:10]
            elif cnt==2:
                data_num=int(line[:7])
            if cnt>data_num:
                box_size=line[:50]
        print(system,data_num,box_size)