#!/usr/bin/python3
import numpy as np
import cv2
import numpy as np
from math import sqrt, atan2
def vector_visualizer(vidIn,CalIn,vidOut):
	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))

	fs = cv2.FileStorage(CalIn, cv2.FILE_STORAGE_READ)
	fn = fs.getNode("Camera_Matrix")

	mtx=fn.mat()
	fn = fs.getNode("Distortion_Coefficients")
	dist=fn.mat()
	print(mtx)
	print(dist)
	while(cap.isOpened() and count<frameCount-1):
		ret, frame = cap.read()
		if ret==True:
			#Jason's Magic
			h,  w = frame.shape[:2]
			print((w,h))
			newcameramtx, roi=cv2.getOptimalNewCameraMatrix(mtx,dist,(w,h),1,(w,h))

			# undistort
			dst = cv2.undistort(frame, mtx, dist, None, newcameramtx)

			# crop the image
			x,y,w,h = roi
			print((x,y,w,h))
			dst = dst[y:y+h, x:x+w]
			cv2.imwrite('testc.png',frame)
			cv2.imwrite('test.png',dst)
			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('CAL',type=str, help="location of calibreation file", default="",nargs='?')
	parser.add_argument('DEST',type=str, help="output video (avi)", default="",nargs='?')
	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.CAL=="":
		inputV=input("location of piv####.txt files (leave blank for default): ")
		if not inputV=="":
			pivData = inputV
	else:
		pivData=args.CAL

	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)