import cv2
import matplotlib.pyplot as plt
import numpy as np
import csv
from scipy.signal import savgol_filter

startind=767
ind=startind
scale=(755-143)/(7*12*25.4/1000)#pixels/m


cap = cv2.VideoCapture("toss_long.mp4")
ret, frame = cap.read()

cap.set(cv2.CAP_PROP_POS_FRAMES, ind)



tracker = cv2.TrackerMIL_create()
bbox=[90,435,30,30]

ret, frame = cap.read()
g = frame[:,:,1]
#initialize tracker, set bounding boxt
ok=tracker.init(g,bbox)

width  = int(cap.get(3))   # float `width`
height = int(cap.get(4))  # float `height`
print([width,height])
output = cv2.VideoWriter('output_video_drop.avi', cv2.VideoWriter_fourcc('M','J','P','G'), 20, [width, height])

xpositions=[]
ypositions=[]
times=[]
smxpositions=[]
smypositions=[]

while(1):
   ret, frame = cap.read()
   if cv2.waitKey(1) & 0xFF == ord('q') or ret==False :
       cap.release()
       cv2.destroyAllWindows()
       break
   #frame[660:800,240:800,:]=0#blackout the mirror 
   g = frame[:,:,1]
   ok,bbox = tracker.update(g) #update tracking
   ypositions.append(-(bbox[1]+bbox[3]/2)/scale)
   xpositions.append((bbox[0]+bbox[2]/2)/scale)
   times.append((ind-startind)/120)#240 fps
   if len(times)>10:
      smxpositions=savgol_filter(xpositions, 10, 3)
      smypositions=savgol_filter(ypositions, 10, 3)
      for i in range(len(xpositions)-1):
         cv2.line(frame, (int(smxpositions[i]*scale),int(-smypositions[i]*scale)),(int(smxpositions[i+1]*scale),int(-smypositions[i+1]*scale)) , (255, 0, 0), 2)#Draw Bounding Box Around Object
         
   font = cv2.FONT_HERSHEY_SIMPLEX
   cv2.putText(frame,'t='+"{:.3f}".format((ind-startind)/120)+' s',(50,150),font, 4,(255,255,255),2,cv2.LINE_AA)
   cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[0]+bbox[2], bbox[1]+bbox[3]), (0, 255, 0),  2)#Draw Bounding Box Around Object   
   for i in range(40):
      cv2.line(frame, (0,145+round(scale*i/10)),(720,145+round(scale*i/10)) , (255, 0, 255), 1)#Draw Bounding Box Around Object
      cv2.line(frame, (120+round(scale*i/10),0),(120+round(scale*i/10),1280) , (255, 0, 255), 1)#Draw Bounding Box Around Object
   #for i in range(len(xpositions)-1):
    #  cv2.line(frame, (int(xpositions[i]*scale),int(-ypositions[i]*scale)),(int(xpositions[i+1]*scale),int(-ypositions[i+1]*scale)) , (255, 0, 0), 2)#Draw Bounding Box Around Object
         
   cv2.imshow('frame',frame)
   b = frame[:,:,0]
   g = frame[:,:,1]
   r = frame[:,:,2]    
   bimg = np.zeros(frame.shape)
   bimg[:,:,0] = b
   bimg=cv2.convertScaleAbs(bimg)
   gimg = np.zeros(frame.shape)
   gimg[:,:,1] = g
   gimg=cv2.convertScaleAbs(gimg)
   rimg = np.zeros(frame.shape)
   rimg[:,:,2] = r
   rimg=cv2.convertScaleAbs(rimg)
   # Display the resulting frame
##   cv2.imshow('blue', bimg)
   cv2.imshow('green', g)
##   cv2.imshow('red', rimg)
   
   #print(frame.size)
   fps = cap.get(cv2.CAP_PROP_FPS)
   #print(fps)
   #print(ind)
   timer=cv2.getTickCount()
   #write video frame to image
   if (ind-startind)%6==0:
      cv2.imwrite('.\\toss_frames\\'+str(ind)+'.jpg', frame)

   #write video frame to video
   output.write(frame)
   ind+=1
   
   if ind>900:#235:
      break

output.release()
cap.release()


np.savetxt('drop.csv', (times, ypositions,xpositions), delimiter=',')


#fit and plot output from object tracking as time series data

thefit = np.polyfit(times,ypositions,2)
plt.subplot(2,1,1)
plt.plot(times,ypositions,'bx-')
plt.plot(times,np.multiply(np.power(times,2),thefit[0])+np.multiply(times,thefit[1])+thefit[2],'r')
plt.xlabel('times (s)')
plt.ylabel('y position (m)')

fit2=np.polyfit(times,(xpositions),1)
plt.subplot(2,1,2)
plt.plot(times,(xpositions),'bx-')
plt.plot(times,(np.multiply(times,fit2[0])+fit2[1]),'r')
#plt.plot(times,np.exp(np.multiply(times,fit2[0])+fit2[1]),'r')
plt.xlabel('times (s)')
plt.ylabel('x position (m)')
print(thefit[0]*2)
plt.show()
