import argparse
import datetime
import json
import time
import cv2
import numpy as np
import os

#Script allows user to extract the coordinates of an ant by clicking on the image taken from a video
#After giving the directory, the name of the python script, and the name of the video, the first image appears
#The user can click on the head of an ant then press the space bar to go 1s later in the video, click again, press space bar, etc.
#Press Q key to end the analysis.
#The Script will create an image of the trajectory, a file with X,Y corrdinates of each click,
#and a file with the calculated distance walked by the ant

# construct the argument parser and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-v", "--video", help="path to the video file")
args = vars(ap.parse_args())

delais_analyse = 15.0
fps = 15.0
time_step = delais_analyse / fps

camera = cv2.VideoCapture(args["video"])
filename, fileext = os.path.splitext(args["video"])

cropping = [0, 1080, 0, 1920]

filename, extension = os.path.splitext(args["video"])

def click_and_store(event, x, y, flags, image):
	if event == cv2.EVENT_LBUTTONDOWN:
		global coords_temp_x, coords_temp_y
		coords_temp_x = x
		coords_temp_y = y


def compute_distance(Coordonnees):
	distance = np.ndarray(shape = (Coordonnees.shape[0]-1,))
	for i in range(1, Coordonnees.shape[0]):
		distance[i-1] = np.sqrt( (Coordonnees[i,0] - Coordonnees[i-1,0])**2 + (Coordonnees[i,1] - Coordonnees[i-1,1])**2 )

	return distance

cv2.namedWindow("Video", cv2.WINDOW_NORMAL)
cv2.setMouseCallback("Video", click_and_store)

Coord_X = []
Coord_Y = []
first_pic = None
frame_ID = 0

# loop over the frames of the video
while True:
	(grabbed, frame) = camera.read()

	if not grabbed:
		break

	frame = frame[cropping[0]:cropping[1], cropping[2]:cropping[3], :]
	
	if first_pic is None:
		first_pic = frame
		
	##########

	cv2.imshow("Video", frame)
	if np.mod(frame_ID, delais_analyse) == 0:
		coords_temp_x = 0
		coords_temp_y = 0
		key = cv2.waitKey(0) & 0xFF
		Coord_X.append(coords_temp_x)
		Coord_Y.append(coords_temp_y)
	else:
		key = cv2.waitKey(1) & 0xFF
		
	#if the 'q' key is pressed, break from the lop
	if key == ord("q"):
		break

	frame_ID += 1

Coordonnees = np.ndarray(shape = (len(Coord_X), 2)).astype('int')
Coordonnees[:,0] = np.array(Coord_X)
Coordonnees[:,1] = np.array(Coord_Y)

cv2.circle(first_pic,(Coordonnees[0,0],Coordonnees[0,1]), 5, (0,0,255), -1)
for i in range(1, Coordonnees.shape[0]):
	cv2.line(first_pic,(Coordonnees[i-1,0],Coordonnees[i-1,1]),(Coordonnees[i,0],Coordonnees[i,1]),(255,0,0), 2)
	cv2.circle(first_pic,(Coordonnees[i,0],Coordonnees[i,1]), 5, (0,0,255), -1)

distance = compute_distance(Coordonnees)

# cleanup the camera and close any open windows

cv2.imwrite(filename + "_Trip.png", first_pic)
np.savetxt(filename + "_Coordonnees.txt", Coordonnees, fmt = '%i')
np.savetxt(filename + "_Distance.txt", distance, fmt = '%1.4f')

camera.release()
cv2.destroyAllWindows()