#!/usr/bin/python3
import numpy as np
import cv2
from math import sqrt, atan2
def vector_visualizer(vidIn,frameIn,vidOut,number=False):
	cap = cv2.VideoCapture(vidIn)

	# Define the codec and create VideoWriter object
	fourcc = cv2.VideoWriter_fourcc(*'XVID')
	frameRate = cap.get(cv2.CAP_PROP_FPS)
	framesize=(round(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),round(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
	out = cv2.VideoWriter(vidOut,fourcc, frameRate, framesize)
	count=0
	frameCount = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

	averages={}
	counts={}
	lastDir=0
	while(cap.isOpened() and count<frameCount-1 and count<3000):
		ret, frame = cap.read()
		if ret==True:
			vectorField=open(frameIn+"/piv"+str(count)+".txt" )
			for line in vectorField:

				parseLine=line.split('\t')
				x=round(float(parseLine[0]))
				y=framesize[1]-round(float(parseLine[1]))
				u=round(float(parseLine[2]))
				v=-round(float(parseLine[3]))
				m=float(parseLine[4])
				#if(abs(
				if(not ((x,y) in averages)):
					averages[(x,y)]=sqrt(u*u+v*v)
					counts[(x,y)]=1
				else:
					tmp=averages[(x,y)]
					averages[(x,y)]=tmp+sqrt(u*u+v*v)
					counts[(x,y)]+=1			
				if m==0:
					m=255
				else:
					m=0
				mag=sqrt(u*u+v*v)
				avg=averages[(x,y)]/counts[(x,y)]
				if mag>avg*2/3 and mag<avg*3:
					if number:
						cv2.putText(frame,str(round((u*u+v*v)**0.5)),(x,y),cv2.FONT_HERSHEY_SIMPLEX,.3,(0,0,0),2)
					else:
						cv2.arrowedLine(frame,(x,y),(x+u,y+v),(m,0,255-m))
						
			out.write(frame)
			count+=1
			print(count)
		else:
			break
	print('\a')
	# Release everything if job is finished
	cap.release()
	out.release()
	cv2.destroyAllWindows()

if __name__ == '__main__':
	import os
	from argparse import *
	cwd=os.getcwd()
	parser=ArgumentParser(description="Tool to overlay velocity vectors on footage")
	parser.add_argument('VIDEO',type=str, help="Background Video", default="",nargs='?')
	parser.add_argument('FILES',type=str, help="location of piv####.txt files", default="",nargs='?')
	parser.add_argument('DEST',type=str, help="output video (avi)", default="",nargs='?')

	parser.add_argument('-n',help="use numbers instead of arrows",action='store_true')
	args = parser.parse_args()
	videoInputPath = cwd+r'/DJI_0001.MOV'
	pivData=cwd+'/OpenPIV_Out'
	videoOutputPath = cwd+r'/output2.avi'

	
	if args.VIDEO=="":
		inputV=input("Video to use (leave blank for default): ")
		if not inputV=="":
			videoInputPath = inputV
	else:
		videoInputPath=args.VIDEO

	if args.FILES=="":
		inputV=input("location of piv####.txt files (leave blank for default): ")
		if not inputV=="":
			pivData = inputV
	else:
		pivData=args.FILES

	if args.DEST=="":
		inputV=input("file to export(leave blank for default): ")
		if not inputV=="":
			videoOutputPath = inputV
	else:
		videoOutputPath=args.DEST	
		
	vector_visualizer(videoInputPath,pivData,videoOutputPath,args.n)