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)
            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 and cnt<=data_num+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])
            elif cnt>data_num+2:
                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 subdomain for each frame of 
    of .trr frame

    parameters:     data = {dictionary input from 'def fr_export'}
                    dx = desired size X of subdomain {units of input}
                    dy = desired size Y of subdomain {units of input}
                    dz = desired size Z of subdomain {units of input}
    
    output:         dictionary of x,y,z grid points per step (frame)
    """
    box_p={}
    for step in data.keys():
        box_p[step]={}
        xs=int(round(data[step]['box'][0][0]+0.1)) #to round always to bigger number
        ys=int(round(data[step]['box'][1][1]+0.1))
        zs=int(round(data[step]['box'][2][2]+0.1))
        box_x=int(xs/dx)
        box_y=int(ys/dy)
        box_z=int(zs/dz)

        xx=[]
        for i in range(0,xs+1,box_x):
            xx.append(i)
        box_p[step]['x']=xx

        yy=[]
        for i in range(0,ys+1,box_y):
            yy.append(i)
        box_p[step]['y']=yy

        zz=[]
        for i in range(0,zs+1,box_z):
            zz.append(i)
        box_p[step]['z']=zz


        xyz=[]
        for ix in range(0,xs+1,box_x):
          for iy in range(0,ys+1,box_y):
            for iz in range(0,zs+1,box_z):
                xyz.append([ix,iy,iz])
        box_p[step]['xyz']=xyz

    return box_p

def resid_data(atom_type, group=[]):
    """
    Finds the index in list that 'def read_gro' returns,
    and correspond to the atom types in group list

    parameters:     atom_type=[]
                    group=[]

    output:         dictionary
    """
    ndx={}
    for element in group:
        ndx[element]=[i for i, e in enumerate(atom_type) if e.strip() == element.strip()]
        #print(element,ndx)
    return ndx

def res2grid(data, atom_num, box_p, ndx):
    """
    Assign residue that corresponds to ndx (see 'def resid_data')
    to sudomains created from 'def domain_decomposition'. The output
    is a the box location (non zero based) in xyz-space for each atom number.

    parameters:     data = {dictionary input from 'def fr_export'}
                    atom_num = [list from 'def read_gro']
                    box_p = {dictionary input from 'def domain_decomposition'}
                    ndx = {dictionary input from 'def resid_data'}

    output:         dictionaries
    """

    box_res={}
    box_res_rev = {}
    for step in data.keys():
        box_res[step]={}
        box_res_rev[step] = {}
        for key in ndx.keys():
            for i in ndx[key]:
                #data[step]['x'][atom_num][x(0),y(1),z(2)]

                xx=data[step]['x'][i][0]
                cnt_x=-1
                prev=-1
                for ix in box_p[step]['x']:
                    cnt_x=cnt_x+1
                    if xx<ix and xx>prev:
                        prev=ix
                        break

                yy=data[step]['x'][i][1]
                cnt_y=-1
                prev=-1
                for iy in box_p[step]['y']:
                    cnt_y=cnt_y+1
                    if yy<iy and yy>prev:
                        prev=iy
                        break

                zz=data[step]['x'][i][2]
                cnt_z=-1
                prev=-1
                for iz in box_p[step]['z']:
                    cnt_z=cnt_z+1
                    if zz<iz and zz>prev:
                        prev=iz
                        break
                box_res[step][atom_num[i]]=(cnt_x,cnt_y,cnt_z)
        #Make subdomain position the key and group the residues
        for key, value in sorted(box_res[step].items()):
            box_res_rev[step].setdefault(value, []).append(key.strip())
    
    return box_res,box_res_rev