import cv2
import sys
import numpy as np
import imutils as im

IMG_SCALE_UP = 1
SEARCH_AREA = 0.2

(major_ver, minor_ver, subminor_ver) = cv2.__version__.split('.')
print(cv2.__version__)

def bbox_extend(bbox: cv2.typing.Rect, search_area: float):
	we2 = int(round(bbox[2]*search_area/2))
	he2 = int(round(bbox[3]*search_area/2))
	res = (int(round(bbox[0]-we2)),)
	res += (int(round(bbox[1]-he2)),)
	res += (bbox[2]+2*we2,)
	res += (bbox[3]+2*he2,)

	return res


def image_crop(src, bbox: cv2.typing.Rect, search_area: float = 0):
	bbox = bbox_extend(bbox, search_area)
	x = bbox[0]
	y = bbox[1]
	w = bbox[2]
	h = bbox[3]

	dst = src[y:y+h, x:x+w]
	return dst


def process_image_bbox(src):
	# convert to gray
	gray = cv2.cvtColor(src, cv2.COLOR_BGR2GRAY)

	return gray


def match(image_bbox, image_template, current_angle):
	best_angle = 0
	best_r = 0
	best_maxLoc = 0

	for dangle in np.linspace(-1.0, +1.0, 10):
		# scale up to IMG_SCALE_UP-size
		image_template_rotated = im.resize(image_template, IMG_SCALE_UP*w, IMG_SCALE_UP*h)
		# rotated
		image_template_rotated = im.rotate(image_template_rotated, current_angle + dangle)
		# scale to original-size
		image_template_rotated = im.resize(image_template_rotated, w, h)

		# Perform template match operations
		res = cv2.matchTemplate(image_bbox, image_template_rotated, cv2.TM_CCOEFF_NORMED)
		(minVal, maxVal, minLoc, maxLoc) = cv2.minMaxLoc(res)

		r = maxVal
		if r > best_r:
			best_r = r
			best_angle = dangle
			best_maxLoc = maxLoc

	current_angle += best_angle

	return current_angle, best_maxLoc


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/production_id 4525346 (1080p).mp4")

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

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

# Uncomment the line below to select a different bounding box
bbox = cv2.selectROI("Frame", frame, False)
image_template = process_image_bbox(image_crop(frame, bbox, 0.5))
w, h = image_template.shape[::-1]

# Initialize tracker with first frame and bounding box
tracker.create()
tracker.init(frame, bbox)

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

	# Start timer
	timer = cv2.getTickCount()

	# Update tracker
	ok, bbox = tracker.update(frame)

	image_bbox = process_image_bbox(image_crop(frame, bbox, SEARCH_AREA))
	image_template_rotated = image_template

	if ok:
		current_angle, best_maxLoc = match(image_bbox, image_template, current_angle)
		image_template_rotated = im.rotate(image_template, current_angle)

		# determine the starting and ending (x, y)-coordinates of the
		# bounding box
		(startX, startY) = best_maxLoc
		startX += bbox[0]
		startY += bbox[1]
		endX = startX + bbox[2]
		endY = startY + bbox[3]

		# draw the bounding box on the image
		cv2.rectangle(frame, (startX, startY), (endX, endY), (255, 0, 0), 3)

		# Calculate Frames per second (FPS)
		fps = cv2.getTickFrequency() / (cv2.getTickCount() - timer)
	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, (50, 170, 50), 2)

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

	# Display result
	cv2.imshow("Tracking", frame)
	cv2.imshow("image_bbox", image_template_rotated)

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