from itertools import islice
import os
import sys
import numpy as np
import pandas as pd
import trackpy as tp
import matplotlib as mpl
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit


##################################
##################################
##################################
##############INPUTS##############
kfib_num = 6 # number of k-fibers, changes based on how many trackable k-fibers there are in the spindle 
scalefactor = 0.0657 # specific to our camera pixel size. 1 pixel = 0.0657 microns

path = "/Users/poojasuresh/Documents/iPQB/dumontdata_analysis/WT/cell001"
# input csv file with stored coordinates post k-fiber tracking in order to parse the coordinates using the pandas package and plot the strain maps 
# file format: frame, particle#, x_coord, y_coord

t = pd.read_csv("%s/coord_data.csv"%path)
framenum = list(range(1,3))
##################################


# tracked k-fiber coordinates in the last frame
x_list_last=[]
y_list_last=[]
for el in t.loc[(t['frame'] == len(framenum))]['x']:
    x_list_last.append(el*scalefactor)
for el in t.loc[(t['frame'] == len(framenum))]['y']:
    y_list_last.append(float(el)*float(scalefactor))

# tracked k-fiber coordinates in the first frame
x_list_1=[]
y_list_1=[]
for el in t.loc[(t['frame'] == 1)]['x']:
    x_list_1.append(el*scalefactor)
for el in t.loc[(t['frame'] == 1)]['y']:
    y_list_1.append(el*scalefactor)


# Binning: This bins together n consecutive coordinate to plot, where bin size = n

binsize = 5
numberofbins = kfib_num*100/binsize # 100 is the number of particles per k-fiber. particle 1-100 = k-fiber #1; particle 101-200 = k-fiber #2 and so on  

bin_list=[]
def multiples(m, count):
    for i in range(int(count)):
        num=i*m
        bin_list.append(int(num))

multiples(binsize,numberofbins)

t['bin'] = pd.cut(t['particle'], bin_list)
t.to_csv("%s/coord_data_with_bins.csv"%path)

binned_particles = []
for i in framenum:
    for j in range(1,len(bin_list)):
        binned_particles.append(j)


df = t.groupby(['frame', 'bin'], as_index=False)['x', 'y'].mean()
df['particle'] = pd.DataFrame({'particles': binned_particles})
df.to_csv("%s/binned_coord_data.csv"%path)



#Plotting
final = pd.read_csv('%s/binned_coord_data.csv'%path)

#last frame
x_list_last=[]
y_list_last=[]
for el in final.loc[(final['frame'] == len(framenum))]['x']:
    x_list_last.append(el*scalefactor)
for el in final.loc[(final['frame'] == len(framenum))]['y']:
    y_list_last.append(el*scalefactor)

# first frame
x_list_1=[]
y_list_1=[]
for el in final.loc[(final['frame'] == 1)]['x']:
    x_list_1.append(el*scalefactor)
for el in final.loc[(final['frame'] == 1)]['y']:
    y_list_1.append(el*scalefactor)


fig=plt.figure()
plt.plot(x_list_last, y_list_last, "g*", ms=10)
plt.plot(x_list_1, y_list_1, "mo", ms=8)
#plt.ylim((3,27))
plt.tick_params(axis='both', which='major', labelsize=20)
tp.plot_traj(final, mpp = scalefactor, plot_style={'color': 'k'})
fig.savefig("%s/strain_map.png"%path)
plt.show()