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'):
    """
    Count total frames of .trr file

    parameters:     trajfile = [.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):
    """
    Export frames from gromacs .trr file. 

    parameters:     trajfile = [.trr]
                    num_frames = [number of frames to keep
                                  counting from the end of file]
    """
    cnt_fr=count_frames(trajfile)
    data_all={}
    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)
        for i in range(cnt_fr-num_frames,cnt_fr):
            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])
            step='{step}'.format(**header)
            data_all[step]=data
        return data_all

def read_gro(gro):
    """
    Read .gro file and exports 
    1. system name
    2. data number
    3. box size
    4. residue number
    5. residue type
    6. atom type
    7. atom number
    8. free format data (x,y,z,v,u,w)

    parameters:     gro = [.gro]
    """
    cnt=0
    data_num=0
    res_num = [] 
    res_type = []  
    atom_type = [] 
    atom_num  = []
    rest_dt = []
    with open(gro, 'r') as F:
        for line in F:
            cnt=cnt+1
            print(cnt)
            if cnt>2:
                res_num.append(line[:5])  
                res_type.append(line[5:10])  
                atom_type.append(line[10:15]) 
                atom_num.append(line[15:20])
                rest_dt.append(line[20:])
            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)
        return system,data_num,box_size,res_num,res_type,atom_type,atom_num,rest_dt

def domain_decomposition(data,dx,dy,dz):
    """
    Identify residues that rely inside a subdomain 
    of a rectangular box

    """
    print('xs ys zs x y z')
    for step in data.keys():
        #print(data[step]['box'])
        #data[step]['box'][row][element]
        xs=data[step]['box'][0][0]
        ys=data[step]['box'][1][1]
        zs=data[step]['box'][2][2]
        #data[step]['x'][atom_num][x(0),y(1),z(3)]
        x=data[step]['x'][0][0]
        y=data[step]['x'][0][1]
        z=data[step]['x'][0][2]
        print(xs,ys,zs, x, y, z)
    #Todo: identify points inside box that belong to same residue
    #return