#!/usr/bin/python3
import os

#Establish file path for directory containing OpenPIV output, count files in directory

def reverse_bathymetry(pivData,average,depth,scale):
	fileCount = 0
	fileList = os.listdir(pivData)
	for fileName in fileList:
	    fileCount = fileCount + 1

	#Create array to hold average velocity values
	videoAvgVelocity = []
	averages={}
	counts={}
	avgVel={}
	velCnt={}
	#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]))
			vel=((u**2)+(v**2))**0.5
			if(not ((x,y) in avgVel)):
				avgVel[(x,y)]=vel
				velCnt[(x,y)]=1
			else:
				avgVel[(x,y)]+=vel
				velCnt[(x,y)]+=1
			if vel>avgVel[(x,y)]/velCnt[(x,y)]:
				
				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()
	bathymetry={}
	averagesFile=open(average,'w')
	depthsFile=open(depth,'w')
	for key in averages.keys():
		vector=averages[key]
		print(vector)
		u=vector[0]/counts[key]
		v=vector[1]/counts[key]
		print(counts[key])
		print((u,v))
		print(scale)
		u=u/scale
		v=v/scale
		averages[key]=(u,v)
		speed = (((u**2)+(v**2))**0.5)
		bathymetry[key]=speed*speed/9.81*8#factor of 8 for correction
	
		averagesFile.write(str(key[0])+'\t'+str(key[1])+'\t'+str(u)+'\t'+str(v)+'\n')#keep og coords 
		depthsFile.write(str(key[0]/scale)+'\t'+str(key[1]/scale)+'\t'+str(bathymetry[key])+'\n')#convert to meters
	averagesFile.close()
	depthsFile.close()

if __name__ == '__main__':
	import os
	from argparse import *
	cwd=os.getcwd()
	parser=ArgumentParser(description="Tool to determin depth from speed")
	parser.add_argument('FILES',type=str, help="location of piv####.txt files", default="",nargs='?')
	parser.add_argument('AVERAGE',type=str, help="location of the file for average velocities", default="",nargs='?')
	parser.add_argument('DEPTH',type=str, help="location of the file for depth at each point", default="",nargs='?')
	parser.add_argument('-s',dest="scale",type=str, help="Pixels per meter")
	args = parser.parse_args()
	
	pivData=cwd+'/OpenPIV_Out'
	averages=cwd+'/averages.txt'
	depth=cwd+'/depth.txt'
	scale=1
	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.AVERAGE=="":
		inputV=input("location to export averages (leave blank for default): ")
		if not inputV=="":
			averages= inputV
	else:
		averages=args.AVERAGE

	if args.DEPTH=="":
		inputV=input("location of export depths (leave blank for default): ")
		if not inputV=="":
			depth = inputV
	else:
		depth=args.DEPTH

	if args.scale:
		scale=float(args.scale)
	else:
		inputV=input("how many pixels per meter? (leave blank for default): ")
		if not inputV=="":
			scale = float(inputV)
	reverse_bathymetry(pivData,averages,depth,scale)