%matplotlib inline
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import pandas as pd
import os
import scipy.optimize as so
os.chdir('/Users/q6600sl/Documents/Work/vial_detector')
from vial_detector import *

D = detector('./examples/Test_data.h264')

print 'Video Dimension:', D.img_dim
print 'Video Framerate:', D.frame_rate

#Show a frame of the video
plt.imshow(D.img_stack[0])

D.show_ROI((500, 740), (370, 540))

D.set_ROI((500, 740), (370, 540))

D.show_30frames()

D.set_backgrd((0,5))

plt.imshow(D.bkgrd_image, cmap=cm.Greys_r)

fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, figsize=(13, 6))
ax1.imshow(D.crtd_image[100], cmap=cm.Greys_r)
ax2.imshow(D.img_stack[100], cmap=cm.Greys_r)
ax1.set_title('Background substracted image')
ax2.set_title('Original image')

df = D.bckgrd_pdetect(thld     = 125, 
                      diameter = 5, 
                      maxsize  = 7,
                      minmass  = 300,
                      invert   = True)

df[1].head() #result is a list of pandas.DataFrame, one for each frame of video

fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, figsize=(13, 6))
D.diagnostic_plot(10, ax=ax1)
D.diagnostic_plot(140, ax=ax2)
ax1.set_title('Frame #11')
ax2.set_title('Frame #141')

big_df = pd.concat(df)
big_df = big_df[big_df.True_particle]
grp_by = big_df.groupby('Timestamp')

fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, figsize=(13, 4))
grp_by.y.mean().plot(ax=ax1)
grp_by.y.count().plot(ax=ax2)
ax1.set_ylabel('Mean vertical position')
ax2.set_ylabel('Count of flies that have moved')

F = lambda t, r, n: n*(1-(1-r)**t)
count_df = grp_by.y.count()
p, vcov  = so.curve_fit(F,
                        count_df.index.values,
                        count_df, p0=(0.25, 25.))
print p
print vcov

_x = np.linspace(0, 6)
plt.plot(_x, F(_x, *p), label='Model Fit')
grp_by.y.count().plot(ax=plt.gca(), label='Observed Data')
plt.ylabel('Count of flies that have moved')
plt.legend(loc=4)

tracks = D.particle_tracking(5, 25, memory=5)

ax = plt.subplot(111)
D.diagnostic_plot(-1, ax=ax)
tp.plot_traj(tracks, ax=ax)
ax.set_xlim(0, D.img_stack[-1].shape[-1])
ax.set_ylim(0, D.img_stack[-1].shape[0])
ax.invert_yaxis()
plt.title('Tracks of the particles')

tracks.head()