#Trackmate turning angle calculation
import statsmodels.api as sm
import numpy as np
import pandas as pd
#import pims
#import trackpy as tp
import os

import matplotlib  as mpl 
import matplotlib.pyplot as plt 
import glob
import os
import cv2
import warnings
warnings.filterwarnings("ignore", category=RuntimeWarning) 
directory='C:/Users/marcb/Desktop/trackmate_bgs/sim_trajectories'
substring='z'
os.chdir('C:/Users/marcb/Desktop/flow_trajectories/tmate_angles')
for filename in os.listdir(directory):
    if substring in filename:
        print ('Working on ' + str(filename))
        timename=filename[:-15]+'_time.xlsx'
        t1 = pd.read_csv(os.path.join(directory,filename),header=3,usecols=[2,4,5,7],names=['particle','x','y','frame'])
        time=pd.read_excel(os.path.join(directory,timename), usecols=[1])
        dt=[]
        time.dropna(inplace=True)
        for count, value in enumerate (time.iloc):
            dt.append(time.iloc[count+1][0]-time.iloc[count][0])
            if (count+2)==len(time):
                break
            if dt[count]<0:
                print ('Time file is erroneous. T[0] = time.iloc[403]', count)
                time.drop(time.index[0:count+1], inplace=True)
                time.reset_index(drop=True, inplace=True)
            else:
                continue
        unique_part=[]
        for i in t1.particle.unique():
            unique_part.append(i)
        unique_part = [x for x in unique_part if str(x) != 'nan']
        traj_stats=[]
        dist=[]
        count=0
        for n in range(len(unique_part)):
            tslice=t1.particle==unique_part[n]
            traj_x=t1[tslice].sort_values(['frame']).x
            traj_y=t1[tslice].sort_values(['frame']).y
            frame=t1[tslice].sort_values(['frame']).frame
            if np.any(frame>len(time)):
                break
            start_ind=time.index[int(frame.iloc[0])]
            current_time=time.iloc[start_ind:start_ind+len(frame)]
            dt=current_time.diff()
            #frame_diff=frame.diff()
            x_diff=traj_x.diff()
            y_diff=traj_y.diff()
            vel_dist=np.sqrt(x_diff**2+y_diff**2)
            vel=np.reshape(vel_dist.values,[len(vel_dist),1])/np.reshape(np.array(dt),[len(vel_dist),1])
 
            angles=np.zeros((len(x_diff),1))
            for i in range(len(x_diff)-1):
                theta_1=np.arctan(y_diff.iloc[i]/x_diff.iloc[i])
                theta_2=np.arctan(y_diff.iloc[i+1]/x_diff.iloc[i+1])
                angles[i]=(theta_2-theta_1)*(180/np.pi)
            #    if angles[i]==0:
            #        angles[i]=np.nan
            #        vel[i]=np.nan
            #    if i==len(x_diff)-1:
            #        vel[i]=np.nan
            #angles[i+1]=np.nan
            for count, val in enumerate (angles):
                if val==0:
                    angles[count]=np.nan
            vel=vel[0:len(angles)]
                    
            traj_stats.append([n, angles,vel])

        angle_array=[]
        vel_array=[]
        for i in range(len(traj_stats)):
            section1=traj_stats[i][1]
            section1=np.reshape(section1,(len(section1),))
            section2=traj_stats[i][2]
            section2=np.reshape(section2,(len(section2),))
            angle_array=np.hstack([section1,angle_array])
            vel_array=np.hstack([section2,vel_array])
            final_array=np.vstack([np.transpose(angle_array)[0:len(vel_array)],np.transpose(vel_array)])

        pd.DataFrame(np.transpose(final_array)).to_csv('vel'+str(filename))