import cv2
import matplotlib.pyplot as plt
import numpy as np
import csv

startind=145
ind=1#startind
scale=(1107-15)/(7*12*25.4/1000)#pixels/m


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

cap.set(cv2.CAP_PROP_POS_FRAMES, ind)



tracker = cv2.TrackerMIL_create()
bbox=[155,0,55,55]

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])

positions=[]
times=[]

while(1):
   ret, frame = cap.read()
   if cv2.waitKey(1) & 0xFF == ord('q') or ret==False :
       cap.release()
       cv2.destroyAllWindows()
       break
   g = frame[:,:,1]
   ok,bbox = tracker.update(g) #update tracking
   positions.append(-(bbox[1]+bbox[3]/2-220)/scale)
   times.append((ind-startind)/120)#240 fps
         
   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(30):
      cv2.line(frame, (150,14+round(scale*i/10)),(250,14+round(scale*i/10)) , (255, 0, 255), 1)#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('.\\drop_frames\\'+str(ind)+'.jpg', frame)

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

output.release()
cap.release()


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


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

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