import cv2
import sys
import imutils as im

from util import bbox_extend, bbox_round, image_crop, to_rect, bbox_add_position, bbox_center

COLOR_TRACKER = (0, 255, 0)
COLOR_TRACKER_EXT = (0, 255, 255)
COLOR_MATCHER = (0, 0, 255)
COLOR_TEMPLATE = (255, 0, 0)
IMG_SCALE_UP = 1
TEMPLATE_MATCH_OVERLAP = 0

video = cv2.VideoCapture("./data/IMG_1730.MP4")

# Exit if video not open
# d.
if not video.isOpened():
	print("Could not open video")
	sys.exit()

# Read first frame.
ok, image = video.read()
if not ok:
	print('Cannot read video file')
	sys.exit()

bb_crop = bbox_extend(cv2.selectROI("Image", image, False), TEMPLATE_MATCH_OVERLAP)
crop = image_crop(image.copy(), bb_crop)

bb_template = cv2.selectROI("Image", crop, fromCenter=True, showCrosshair=True)
template = image_crop(crop.copy(), bb_template)

# Initialize tracker with first frame and bounding box
tracker = cv2.TrackerKCF().create()
tracker.init(image, bb_crop)

path = []
line_from = None

is_startup = True
offset_x = 0
offset_y = 1
while True:
	# Read a new frame
	ok, image = video.read()
	if not ok:
		break

	# Start timer
	timer = cv2.getTickCount()

	# Update tracker
	ok, bb_crop = tracker.update(image)
	crop = image_crop(image.copy(), bb_crop)
	image_anno = image.copy()

	if ok:
		# Perform template match operations
		crop_scaled = im.resize(crop.copy(), IMG_SCALE_UP * crop.shape[1], IMG_SCALE_UP * crop.shape[0])
		template_scaled = im.resize(template, IMG_SCALE_UP * template.shape[0], IMG_SCALE_UP * template.shape[1])
		res = cv2.matchTemplate(cv2.cvtColor(crop_scaled, cv2.COLOR_BGR2GRAY), cv2.cvtColor(template_scaled, cv2.COLOR_BGR2GRAY), cv2.TM_CCOEFF_NORMED)
		(min_val, max_val, min_loc, max_loc) = cv2.minMaxLoc(res)
		print(f"(min_val, max_val, min_loc, max_loc): {(min_val, max_val, min_loc, max_loc)}")
		print(f"matcher_res     : {res.shape}")
		print(f"crop     		: {crop.shape}")
		print(f"crop_scaled     : {crop_scaled.shape}")
		print(f"template        : {template.shape}")
		print(f"template_scaled : {template_scaled.shape}")

		matcher_start_x = max_loc[0]/IMG_SCALE_UP
		matcher_start_y = max_loc[0]/IMG_SCALE_UP
		if is_startup:
			is_startup = False
			offset_x = matcher_start_x - bb_template[0]
			offset_y = matcher_start_y - bb_template[1]
			print(f"offset : {(offset_x, offset_y)}")

		matcher_start_x -= offset_x
		matcher_start_y -= offset_y
		print(f"matcher_start : {(matcher_start_x, matcher_start_y)}")
		matcher_bbox_local = (round(matcher_start_x), round(matcher_start_y), template.shape[1], template.shape[0])
		matcher_top_left_local, matcher_bottom_right_local = to_rect(bbox_round(matcher_bbox_local))

		matcher_bbox = bbox_round(bbox_add_position(matcher_bbox_local, bb_crop))
		matcher_rect = to_rect(matcher_bbox)

		tracker_rect = to_rect(bb_crop)
		template_top_left, template_bottom_right = to_rect(bb_template)

		cv2.rectangle(crop, template_top_left, template_bottom_right, COLOR_TEMPLATE, 1)
		cv2.rectangle(crop, matcher_top_left_local, matcher_bottom_right_local, COLOR_MATCHER, 1)

		cv2.rectangle(image_anno, matcher_rect[0], matcher_rect[1], COLOR_MATCHER, 1)
		cv2.rectangle(image_anno, tracker_rect[0], tracker_rect[1], COLOR_TRACKER, 1)
		matcher_crop = image_crop(image, matcher_bbox)

		# Draw path
		line_to = bbox_center(matcher_bbox)
		if line_from is not None:
			path.append({'from': line_from, 'to': line_to})
		line_from = line_to

		for p in path:
			cv2.line(image_anno, p['from'], p['to'], (0, 0, 255), 2)

		cv2.imshow("Crop", crop)
		cv2.imshow("Image", image_anno)
		cv2.imshow("Matcher res", res)
		cv2.imshow("Matcher view", matcher_crop)

	while True:
		# Exit if ESC pressed
		k = cv2.waitKey(1) & 0xff
		if k == 27:
			break
		if k == 32:
			break

	if k == 27:
		break