#!/usr/bin/python3
import os
import cv2
from depthFrame import depthFrame
#Establish file path for directory containing OpenPIV output, count files in directory

def video_bathymetry(pivData,vidOut,frameRate,framesize,scale):
	fileCount = 0
	fileList = os.listdir(pivData)
	for fileName in fileList:
	    fileCount = fileCount + 1

	#Create array to hold average velocity values
	videoAvgVelocity = []
	averages={}
	counts={}
	fourcc = cv2.VideoWriter_fourcc(*'XVID')
	bathymetry={}
	out = cv2.VideoWriter(vidOut,fourcc, frameRate, framesize)
	#Begin loop to loop through file directory
	for i in range(0, fileCount - 1):
		#Find and open file
		filePath = pivData + "/piv"+str(i)+".txt"
		inFile = open(filePath)

		#Read through file to obtain max, then return to beginning
		rawVectors = inFile.readlines()


		#Read the file data line by line

		for j in rawVectors:
			line = j
			lineInfo = line.split("\t")
			x = float(lineInfo[0])
			y = float(lineInfo[1])
			u = float(lineInfo[2])
			v = float(lineInfo[3])
			mask = round(float(lineInfo[4]))
			if mask<2:
				
				if(not ((x,y) in averages)):
					averages[(x,y)]=(u,v)
					counts[(x,y)]=1
				else:
					tmp=averages[(x,y)]
					averages[(x,y)]=(tmp[0]+u,tmp[1]+v)
					counts[(x,y)]+=1			
				#Calculate speed for each line and append to array
			#speed = ((u**2)+(v**2))**0.5
			
		print(i)
		inFile.close()


		for key in averages.keys():
			vector=averages[key]
			u=vector[0]/counts[key]
			v=vector[1]/counts[key]
			speed = (((u**2)+(v**2))**0.5)/scale
			bathymetry[key]=speed*speed/9.81*81#there was a bug, I had to fix, so I corrected for fixing the bug with a hack, its up to you
		frame=depthFrame(bathymetry,16,framesize[0],framesize[1])
		cv2.imwrite("lastFrame.png",frame)
		#newFrame= cv2.resize(frame,framesize,fx=0,fy=0, interpolation = cv2.INTER_CUBIC)
		out.write(frame)
	out.release()
if __name__ == '__main__':
	import os
	from argparse import *
	cwd=os.getcwd()
	parser=ArgumentParser(description="Tool to make bathimetry video")
	parser.add_argument('FILES',type=str, help="location of piv####.txt files", default="",nargs='?')
	parser.add_argument('VIDEO',type=str, help="location to export video to", default="",nargs='?')
	args = parser.parse_args()
	
	pivData=cwd+'/OpenPIV_Out'
	video="output.avi"
	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.VIDEO=="":
		inputV=input("location to export video to:  ")
		if not inputV=="":
			video= inputV
	else:
		video=args.VIDEO

	
	video_bathymetry(pivData,video,60,(128*2,68*2),58)