import cv2
import sys
import imutils as im

from template_matcher import template_match_rotation
from util import bbox_extend, bbox_round, process_image_bbox, image_crop, bbox_add_position, to_rect

IMG_SCALE_UP = 1
TEMPLATE_MATCH_OVERLAP = 1


def match(source, template, angle, center):
	# coarse
	angle_c, loc, res = template_match_rotation(source, template, angle, center)
	# fine
	angle, loc, res = template_match_rotation(source, template, angle_c, center, rot_min=-0.25, rot_max=+0.25, n_steps=20, scale=IMG_SCALE_UP)

	return angle, loc, res


class Reference:
	def __init__(self, ref_gray):
		self.center = (0, 0)
		self.ref = cv2.cvtColor(ref_gray, cv2.COLOR_GRAY2BGR)
		self.work = self.ref.copy()
		cv2.namedWindow('Ref')
		cv2.setMouseCallback('Ref', self.extract_coordinates)

	def set(self):
		while True:
			cv2.imshow("Ref", self.work)
			key = cv2.waitKey(1)
			if key == 32:
				break
			elif key == 27:
				break
		cv2.destroyWindow("Ref")
		return self.center

	def extract_coordinates(self, event, x, y, flags, parameters):
		if event == cv2.EVENT_LBUTTONDOWN:
			self.work = self.ref.copy()
			self.center = (x, y)

			print(f"Pixel at {self.center} = {self.work[y, x]}")
			self.work[y, x] = (0, 0, 255)
			print(f"Pixel at {self.center} = {self.work[y, x]}")


trackers = {
	'KCF': cv2.TrackerKCF,
	'CSRT': cv2.TrackerCSRT,
	'MIL': cv2.TrackerMIL,
	'GOTURN': cv2.TrackerGOTURN,
	'NANO': cv2.TrackerNano,
	'MOSSE': cv2.legacy.TrackerMOSSE,
	'TLD': cv2.legacy.TrackerTLD,
	'BOOSTING': cv2.legacy.TrackerBoosting,
	'MEDIANFLOW': cv2.legacy.TrackerMedianFlow
}
tracker_type = 'KCF'
tracker = trackers[tracker_type].create()

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, frame_init = video.read()
if not ok:
	print('Cannot read video file')
	sys.exit()

bbox_template = cv2.selectROI("Frame", frame_init, False)
image_template = process_image_bbox(image_crop(frame_init, bbox_template))
ref = Reference(image_template)
ref_center = ref.set()

print(f"ref_center = {ref_center}")
# Initialize tracker with first frame and bounding box
tracker.create()
tracker.init(frame_init, bbox_template)

COLOR_TRACKER = (0, 255, 0)
COLOR_TRACKER_EXT = (0, 255, 255)
COLOR_MATCHER = (0, 0, 255)
current_angle = 0
path = []
line_from = None

while True:
	# Read a new frame
	ok, frame = video.read()
	if not ok:
		break

	# Start timer
	timer = cv2.getTickCount()

	# Update tracker
	ok, bbox_tracker = tracker.update(frame)
	print(f"Bbox Tracker: {bbox_tracker}")
	image_template_rotated = image_template
	bbox_tracker_ext = bbox_extend(bbox_tracker, TEMPLATE_MATCH_OVERLAP)

	if ok:
		crop = image_crop(frame.copy(), bbox_tracker_ext)
		current_angle, max_loc, matcher_res = match(process_image_bbox(crop), image_template, current_angle, ref_center)

		# determine the starting and ending (x, y)-coordinates of the
		# bounding box
		image_template_rotated = im.rotate(image_template, current_angle, center=ref_center)

		# Matcher
		matcher_bbox_local = (max_loc[0], max_loc[1], image_template.shape[1], image_template.shape[0])
		matcher_rect_local = to_rect(bbox_round(matcher_bbox_local))

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

		# Tracker
		tracker_rect = to_rect(bbox_round(bbox_tracker))

		# Tracker ext
		tracker_ext_rect = to_rect(bbox_round(bbox_tracker_ext))

		cv2.rectangle(frame, tracker_ext_rect[0], tracker_ext_rect[1], COLOR_TRACKER_EXT, 2)
		cv2.rectangle(frame, tracker_rect[0], tracker_rect[1], COLOR_TRACKER, 2)
		cv2.rectangle(frame, matcher_rect[0], matcher_rect[1], COLOR_MATCHER, 1)

		# Calculate Frames per second (FPS)
		fps = cv2.getTickFrequency() / (cv2.getTickCount() - timer)

		# Draw path
		line_to = matcher_rect[0]
		if line_from is not None:
			path.append({'from': line_from, 'to': line_to})
		line_from = line_to

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

	else:
		# Tracking failure
		cv2.putText(frame, "Tracking failure detected", (100, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)

	# Display tracker type on frame
	cv2.putText(frame, tracker_type + " Tracker", (100, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.75, COLOR_TRACKER, 2)
	cv2.putText(frame, tracker_type + " Tracker Ext", (100, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.75, COLOR_TRACKER_EXT, 2)
	cv2.putText(frame, tracker_type + " Matcher", (100, 100), cv2.FONT_HERSHEY_SIMPLEX, 0.75, COLOR_MATCHER, 2)

	# Display FPS on frame
	cv2.putText(frame, "FPS : " + str(int(fps)), (100, 140), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (50, 170, 50), 2)

	# Display result
	cv2.imshow("Frame", frame)
	cv2.imshow("Matcher", matcher_res)
	cv2.imshow("Crop", image_crop(frame.copy(), matcher_bbox))

	# Exit if ESC pressed
	k = cv2.waitKey(100) & 0xff
	if k == 27:
		break