# bb-tool.py
#
# This is a small app to aid in extracting bounding boxes from video using a human in the loop and some cv algs
# It builds heavily on the examples from the cvui library: https://github.com/Dovyski/cvui


import cv2
import numpy as np
import os
import sys
from EnhancedWindow import EnhancedWindow
import cvui
import time


# Some globals
VERSION = '1.0.0'
VIDEO_WINDOW_NAME = 'BB-Tool ' + VERSION + ' : Video'
CTRL_WINDOW_NAME = 'BB-Tool ' + VERSION + ' : Control'

TRACKING_PARAMS_FILE = 'tracking_params.json'


def vid_frame(cap, frame):
	# get a new frame from the video
	ret, new_frame = cap.read()
	if not ret:
		return False, []
	new_frame = cv2.resize(new_frame, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_LINEAR)
	return True, new_frame

# Draw and update ROIs in the video (compied almost exactly from complex-mouse cvui example


def roi_tool(frame,anchor,roi):
	# The function 'bool cvui.mouse(int query)' allows you to query the mouse for events.
	# E.g. cvui.mouse(cvui.DOWN)
	#
	# Available queries:
	#	- cvui.DOWN: any mouse button was pressed. cvui.mouse() returns true for single frame only.
	#	- cvui.UP: any mouse button was released. cvui.mouse() returns true for single frame only.
	#	- cvui.CLICK: any mouse button was clicked (went down then up, no matter the amount of frames in between). cvui.mouse() returns true for single frame only.
	#	- cvui.IS_DOWN: any mouse button is currently pressed. cvui.mouse() returns true for as long as the button is down/pressed.

	drawing = False
	new_bb = False

	# Did any mouse button go down?
	if cvui.mouse(cvui.DOWN):
		# Position the anchor at the mouse pointer.
		anchor.x = cvui.mouse().x
		anchor.y = cvui.mouse().y

		# Inform we are working, so the ROI window is not updated every frame
		drawing = True

	# Is any mouse button down (pressed)?
	if cvui.mouse(cvui.IS_DOWN):
		# Adjust roi dimensions according to mouse pointer
		width = cvui.mouse().x - anchor.x
		height = cvui.mouse().y - anchor.y

		roi.x = anchor.x + width if width < 0 else anchor.x
		roi.y = anchor.y + height if height < 0 else anchor.y
		roi.width = abs(width)
		roi.height = abs(height)

		# Show the roi coordinates and size
		cvui.printf(frame, roi.x + 5, roi.y + 5, 0.3, 0xff0000, '(%d,%d)', roi.x, roi.y)
		cvui.printf(frame, cvui.mouse().x + 5, cvui.mouse().y + 5, 0.3, 0xff0000, 'w:%d, h:%d', roi.width, roi.height)

	# Was the mouse clicked (any button went down then up)?
	if cvui.mouse(cvui.UP):
		# We are done working with the ROI.
		drawing = False
		if roi.width > 25 and roi.height > 25:
			new_bb = True

	# Ensure ROI is within bounds
	frameRows, frameCols, frameChannels = frame.shape
	roi.x = 0 if roi.x < 0 else roi.x
	roi.y = 0 if roi.y < 0 else roi.y
	roi.width = roi.width + frame.cols - (roi.x + roi.width) if roi.x + roi.width > frameCols else roi.width
	roi.height = roi.height + frame.rows - (roi.y + roi.height) if roi.y + roi.height > frameRows else roi.height

	# Render the roi
	cvui.rect(frame, roi.x, roi.y, roi.width, roi.height, 0xff0000)

	return anchor, roi, drawing, new_bb



def apply_canny(frame,high,low,kernel=3):

	frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
	frame = cv2.Canny(frame, low, high, kernel)
	frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
	return frame

# the main app

def main():

	frame = np.zeros((540, 960, 3), np.uint8)
	org_frame = np.zeros((540, 960, 3), np.uint8)
	ctrl_frame = np.zeros((540, 300, 3), np.uint8)
	roi = cvui.Rect(0, 0, 0, 0)
	anchor = cvui.Point()
	low_threshold = [50]
	high_threshold = [150]
	playback_fps = [30]
	use_canny = [False]
	detect_rois = [False]
	fuse_track_and_bbox = [False]
	is_paused = False

	# Create a settings window using the EnhancedWindow class.
	settings = EnhancedWindow(10, 50, 270, 270, 'Settings')
	# Create a settings window using the EnhancedWindow class.
	control = EnhancedWindow(10, 300, 270, 100, 'Control')

	tracker = cv2.TrackerKCF_create()
	tracker_init = False
	if os.path.exists(TRACKING_PARAMS_FILE):
		fs = cv2.FileStorage(TRACKING_PARAMS_FILE, cv2.FileStorage_READ)
		tracker.read(fs.getFirstTopLevelNode())

	# Init cvui and tell it to create a OpenCV window, i.e. cv2.namedWindow(WINDOW_NAME).
	cvui.init(CTRL_WINDOW_NAME)
	cvui.init(VIDEO_WINDOW_NAME)



	# Load the video or die
	try:
		cap = cv2.VideoCapture(sys.argv[1])
	except:
		print ('Could not load the video file, will now exit.')
		exit()

	timer = 0.0 # force a new frame on first iteration

	while True:

		# get and prep new frame or quite if end of video

		get_new_frame = time.time()-timer > 1/playback_fps[0]
		if not is_paused and get_new_frame:
			ret, new_frame = vid_frame(cap, frame)
			timer = time.time()
			if not ret:
				break

		# copy the frame
		org_frame[:] = new_frame[:]

		# Should we apply Canny edge?
		if use_canny[0]:
			frame = apply_canny(new_frame,low_threshold[0], high_threshold[0], 3)
		else: 
			# No, so just copy the original image to the displaying frame.
			frame[:] = new_frame[:]

		# Should we detect rois
		if detect_rois[0]:

			if not use_canny[0]:
				frame = apply_canny(new_frame, low_threshold[0], high_threshold[0], 3)
			frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
			kernel = np.ones((5, 5), np.uint8)
			frame = cv2.dilate(frame, kernel, iterations=1)
			frame = cv2.erode(frame, kernel, iterations=1)
			contours, hierarchy = cv2.findContours(frame,
								cv2.RETR_EXTERNAL,
								cv2.CHAIN_APPROX_SIMPLE
			)

			bounding_rects = []
			for indx,c in enumerate(contours):
				if cv2.contourArea(c) > 250:
					br = cv2.boundingRect(c)
					cv2.rectangle(org_frame,br,(0, 255, 0), 3)
					bounding_rects.append(br)
					#cv2.drawContours(org_frame, contours, indx, (0, 255, 0), 3)

			frame = org_frame

		cvui.context(VIDEO_WINDOW_NAME)
		if not tracker_init or is_paused:
			anchor, roi, drawing, new_bb = roi_tool(frame, anchor, roi)

		# if there is a new bounding box and we are not tracking one,
		# init a new tracker
		if new_bb and (not tracker_init or is_paused):
			# Initialize tracker with first frame and bounding box
			tracker = cv2.TrackerKCF_create()
			ok = tracker.init(frame, (roi.x, roi.y, roi.width, roi.height))
			tracker_init = True

		# if we are not actively drawing a bounding box and the tracker is initialized
		# update it with the latest frame from the video
		if not drawing and tracker_init:
			# Update tracker
			ok = True
			if not is_paused:
				ok, roi = tracker.update(frame)
				match_index = 0
				min_diff = 1000000
				if fuse_track_and_bbox[0]:
					for br_ind, br in enumerate(bounding_rects):
						#print(br)
						sum_squared_diff = 0
						for ind in range(0, 4):
							sum_squared_diff += (roi[ind]-br[ind])**2
						sum_squared_diff = np.sqrt(sum_squared_diff)
						if min_diff > sum_squared_diff:
							match_index = br_ind
							min_diff = sum_squared_diff
						#print(min_diff)
					if min_diff < 100:
						roi = bounding_rects[match_index]
						#print(roi)

				roi = cvui.Rect(roi[0], roi[1], roi[2], roi[3])
			if ok or fuse_track_and_bbox:
				cvui.rect(frame, roi.x, roi.y, roi.width, roi.height, 0xff0000)
			else:
				tracker_init = False

		# This function must be called *AFTER* all UI components. It does
		# all the behind the scenes magic to handle mouse clicks, etc.
		cvui.update(VIDEO_WINDOW_NAME)
		# Show everything on the screen
		cv2.imshow(VIDEO_WINDOW_NAME, frame)

		cvui.context(CTRL_WINDOW_NAME)
		# Render the settings window and its content, if it is not minimized.
		settings.begin(ctrl_frame)
		if not settings.isMinimized():
			cvui.checkbox('Use Canny Edge', use_canny)
			cvui.checkbox('Detect ROIs', detect_rois)
			cvui.checkbox('Fuse Track and Bounding Box', fuse_track_and_bbox)
			cvui.trackbar(settings.width() - 20, low_threshold, 5, 150)
			cvui.trackbar(settings.width() - 20, high_threshold, 80, 300)
			cvui.trackbar(settings.width() - 20, playback_fps, 1, 120)
			cvui.space(20) # add 20px of empty space
		settings.end()

		control.begin(ctrl_frame)
		if not control.isMinimized():
			# Should we quit?
			if cvui.button(settings.width() - 20, 20, "Quit"):
				break

			if is_paused:
				if cvui.button(settings.width() - 20, 20, "Resume"):
					is_paused = False
			else:
				if cvui.button(settings.width() - 20, 20, "Pause"):
					is_paused = True

		control.end()


		cvui.update(CTRL_WINDOW_NAME)
		cv2.imshow(CTRL_WINDOW_NAME, ctrl_frame)

		# Check if ESC key was pressed
		if cv2.waitKey(5) == 27:
			break

	# save the tracking params at the end of the loop
	tracker.save(TRACKING_PARAMS_FILE)

if __name__ == '__main__':
    main()