import cv2
import numpy as np
import argparse
import sys
import csv
import json
import matplotlib.pyplot as plt

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

IMG_SCALE_UP = 1
IMG_ROTATE = 0

COLOR_TRACKER = (0, 255, 0)
COLOR_TRACKER_EXT = (0, 255, 255)
COLOR_MATCHER = (0, 0, 255)
COLOR_TEMPLATE = (255, 0, 0)
COLOR_TRACK = (255, 255, 0)
COLOR_CIRCLE = (0, 255, 0)

CONSOLE_DEBUG = False


class Corner:
	def __init__(self, _ref_img: np.array, _ref_bb: np.array, name: str = 'Corner'):
		self.ref_img = _ref_img
		self.ref_bb = _ref_bb
		self.ref_offset = (0, 0)
		self.name = name
		self._path = []
		self.line_from = None
		self.first_frame = True

	def _debug(self, _image, matcher_res, matcher_start_x, matcher_start_y):
		(min_val, max_val, min_loc, max_loc) = cv2.minMaxLoc(matcher_res)

		self._print(f"(min_val, max_val, min_loc, max_loc): {(min_val, max_val, min_loc, max_loc)}")
		self._print(f"matcher_start   : {(matcher_start_x, matcher_start_y)}")
		self._print(f"matcher_res     : {matcher_res.shape}")
		self._print(f"template        : {self.ref_img.shape}")
		self._print(f"offset 		  : {self.ref_offset}")

	def _print(self, s):
		if CONSOLE_DEBUG:
			print(f"{self.name}: {s}")

	def process(self, _image: np.array):
		# Match
		matcher_res = cv2.matchTemplate(_image, self.ref_img, cv2.TM_CCOEFF_NORMED)
		(min_val, max_val, min_loc, max_loc) = cv2.minMaxLoc(matcher_res)

		# Get matcher coord
		matcher_start_x = max_loc[0]/IMG_SCALE_UP
		matcher_start_y = max_loc[1]/IMG_SCALE_UP

		# Store offset at first frame
		if self.first_frame:
			self.first_frame = False
			offset_x = matcher_start_x - self.ref_bb[0]
			offset_y = matcher_start_y - self.ref_bb[1]
			self.ref_offset = (offset_x, offset_y)

		# Compensate offset
		matcher_start_x -= self.ref_offset[0]
		matcher_start_y -= self.ref_offset[1]

		matcher_bbox_local = (matcher_start_x, matcher_start_y, self.ref_img.shape[1], self.ref_img.shape[0])

		self._debug(_image, matcher_res, matcher_start_x, matcher_start_y)

		return matcher_bbox_local, matcher_res

	def path_add(self, point):
		if len(self.path) > 200:
			return
		if self.line_from is not None:
			self._path.append({'from': self.line_from, 'to': point})
		self.line_from = point

	@property
	def path(self):
		return self._path


class CornerTrackerParams:
	def __init__(self, _scale: float = 1.0, _pre_track: bool = False, _show_path: bool = False):
		self.scale = _scale
		self.pre_track = _pre_track
		self.show_path = _show_path

	def from_dict(self, params: dict):
		self.__dict__.update(params)
		return self


class CornerTrackerSettings:
	def __init__(self, _mask_bb=None, _tracking_bb=None):
		self.mask_bb = _mask_bb
		self.tracking_bb = _tracking_bb
		self.matching_tpl_bb = []

	def from_dict(self, settings: dict):
		self.__dict__.update(settings)
		return self

	def append(self, _bb: np.array):
		self.matching_tpl_bb.append(_bb)


class CornerTracker:
	def __init__(self, _params: CornerTrackerParams, color=(0, 255, 0), name: str = 'CornerTracker'):
		self.params = _params
		self.color = color
		self.name = name
		self.settings = None
		self.tracking_bb = None
		self.tracking_mask = None
		self.tracker = None
		self.corner_ref = None
		self.corner_matcher_list = []

	def _print(self, s):
		if CONSOLE_DEBUG:
			print(f"{self.name}: {s}")

	def _debug(self, _image_anno, matcher_bbox_local):
		matcher_bbox = bbox_round(bbox_add_position(matcher_bbox_local, self.tracking_bb))

		matcher_rect = to_rect(matcher_bbox)
		cv2.rectangle(_image_anno, matcher_rect[0], matcher_rect[1], COLOR_MATCHER, 1)

		tracker_rect = to_rect(bbox_round(self.tracking_bb))
		cv2.rectangle(_image_anno, tracker_rect[0], tracker_rect[1], self.color, 1)

	@staticmethod
	def mask_init(_image: np.array):
		# Create mask
		_mask = np.zeros((_image.shape[0], _image.shape[1]), dtype=np.uint8)
		return _mask

	@staticmethod
	def mask_draw(_image: np.array, _mask: np.array):
		print(f"Draw mask object")
		bb = cv2.selectROI("Tracker Reference", _image, False)
		cv2.destroyWindow("Tracker Reference")
		if bbox_center(bb) == (0, 0):
			bb = (0, 0, _image.shape[1], _image.shape[0])

		return bb

	@staticmethod
	def mask_apply(_image: np.array, _mask: np.array):
		_image[_mask == 0] = 0
		_image[_mask != 0] = _image[_mask != 0]

	@staticmethod
	def image_process(_image: np.array):
		_result = cv2.cvtColor(_image, cv2.COLOR_BGR2GRAY)
		'''
		_result = cv2.GaussianBlur(_result, (9, 9), 0)
		_result = cv2.medianBlur(_result, 3)
		'''
		return _result

	def create_from_settings(self, _image: np.array, settings: CornerTrackerSettings):
		_mask = CornerTracker.create_mask(_image, settings.mask_bb)
		image_processed_local, _ = self.create_tracking(_image, _mask, settings.tracking_bb)
		self.create_matching(image_processed_local, settings.tracking_bb, settings.matching_tpl_bb)

		return self.create_matching(image_processed_local, settings.tracking_bb, settings.matching_tpl_bb)

	@staticmethod
	def create_mask(_image: np.array, mask_bb: np.array):
		# Init mask
		_mask = CornerTracker.mask_init(_image)
		rect_upper_left = (mask_bb[0], mask_bb[1])
		rect_lower_right = (mask_bb[0] + mask_bb[2], mask_bb[1] + mask_bb[3])
		# Mask cut out
		cv2.rectangle(_mask, rect_upper_left, rect_lower_right, 255, -1)

		return _mask

	def create_tracking(self, _image: np.array, _mask: np.array, tracking_bb: np.array):
		# Init search area
		tracking_img = image_crop(_image, tracking_bb)
		tracking_mask = image_crop(_mask, tracking_bb)
		CornerTracker.mask_apply(tracking_img, tracking_mask)
		image_processed_local = CornerTracker.image_process(tracking_img)

		# Initialize tracker with first frame and bounding box
		if self.params.pre_track:
			self.tracker = cv2.TrackerKCF.create()
			self.tracker.init(_image, tracking_bb)

		self.tracking_bb = tracking_bb
		self.tracking_mask = image_crop(_mask, tracking_bb)

		return image_processed_local, tracking_img

	def create_matching(self, image_processed_local: np.array, tracking_bb: np.array, matching_tpl_bb_list: list):
		# Create corner matcher
		self.corner_matcher_list = []
		corner_list = []
		count = 1
		for matching_tpl_bb in matching_tpl_bb_list:
			if bbox_center(matching_tpl_bb) == (0, 0):
				break
			corner_list.append(bbox_center(matching_tpl_bb))
			matching_tpl_img = image_crop(image_processed_local, matching_tpl_bb)
			self.corner_matcher_list.append(Corner(matching_tpl_img, matching_tpl_bb, name=f"Corner-{count}"))

			print(f"Corner {count} added")
			count += 1

		if len(corner_list) == 0:
			return False

		# Refine initial corners and store them as reference
		corners_local = self._corner_refine(image_processed_local, corners=corner_list)

		# Transform corners to global
		self.corner_ref = []
		for corner_local in corners_local:
			corner = (corner_local[0] + tracking_bb[0], corner_local[1] + tracking_bb[1])
			self.corner_ref.append(corner)

		return True

	def init_reference_frame(self, _image: np.array):
		settings = CornerTrackerSettings()
		# Draw mask
		_mask = CornerTracker.mask_init(_image)

		# Init mask
		print(f"Draw mask object")
		_mask_bb = cv2.selectROI("Tracker Reference", _image, False)
		cv2.destroyWindow("Tracker Reference")
		if bbox_center(_mask_bb) == (0, 0):
			_mask_bb = (0, 0, _image.shape[1], _image.shape[0])

		settings.mask_bb = _mask_bb
		_mask = self.create_mask(_image, _mask_bb)

		# Select search area
		print(f"Select tracking object")
		tracking_bb = cv2.selectROI("Tracker Reference", _image, False)
		cv2.destroyWindow("Tracker Reference")
		if bbox_center(tracking_bb) == (0, 0):
			return None

		settings.tracking_bb = tracking_bb
		image_processed_local, tracking_img = self.create_tracking(_image, _mask, tracking_bb)

		count = 1
		while True:
			print(f"Draw Corner #{count}")
			matching_tpl_bb = cv2.selectROI("Matcher Reference", tracking_img, fromCenter=True, showCrosshair=True)
			if bbox_center(matching_tpl_bb) == (0, 0):
				break

			settings.matching_tpl_bb.append(matching_tpl_bb)
			print(f"Corner {count} added")
			count += 1

		cv2.destroyWindow("Matcher Reference")
		print(f"Added {count} corners")

		if self.create_matching(image_processed_local, settings.tracking_bb, settings.matching_tpl_bb):
			return settings

	def process(self, _image: np.array, _image_anno: np.array):
		# Update tracker bb
		if self.tracker is not None:
			_ok, self.tracking_bb = self.tracker.update(_image)
			if not _ok:
				return None

		_image_local = image_crop(_image, self.tracking_bb)
		CornerTracker.mask_apply(_image_local, self.tracking_mask)
		_image_processed_local = CornerTracker.image_process(_image_local)

		return self._match(_image_processed_local, _image_anno)

	def _match(self, _image: np.array, _image_anno: np.array):
		corner_list = []
		for _ct in self.corner_matcher_list:
			matcher_bbox_local, matcher = _ct.process(_image)

			# Draw path
			corner = bbox_center(matcher_bbox_local)
			corner_list.append(corner)

			self._debug(_image_anno, matcher_bbox_local)

		# refine corners
		corners_local = self._corner_refine(_image, corner_list)

		# Transform corners to global
		corners_refined = []
		for corner_local in corners_local:
			corner = (corner_local[0] + self.tracking_bb[0], corner_local[1] + self.tracking_bb[1])
			corners_refined.append(corner)

		corners_refined = np.array(corners_refined)

		if self.params.show_path:
			# Create path from global refined corners
			_i = 0
			for corner in corners_refined:
				_ct = self.corner_matcher_list[_i]
				_ct.path_add(corner)
				_i += 1

			# draw path
			for _ct in self.corner_matcher_list:
				for p in _ct.path:
					cv2.line(_image_anno, bbox_round(p['from']), bbox_round(p['to']), COLOR_TRACK, 1)

		distances = []
		for _i in range(0, corners_refined.shape[0]):
			distances.append((corners_refined[_i] - self.corner_ref[_i]))

		_mean_distance = np.mean(distances, axis=0)

		# show refined corners
		for _i in range(0, corners_refined.shape[0]):
			self._print(f" -- Corner Reference [{i}] {self.corner_ref[_i]}")
			self._print(f" -- Corner coarse    [{i}] {corner_list[_i]}")
			self._print(f" -- Corner fine      [{i}] {corners_refined[_i]}")
			self._print(f" -- Corner distance  [{i}] {corners_refined[_i] - self.corner_ref[_i]}")
			cv2.circle(_image_anno, (int(corners_refined[_i, 0] + 0.5), int(corners_refined[_i, 1] + 0.5)), 4, COLOR_CIRCLE)

		return _mean_distance

	@staticmethod
	def _corner_refine(src_gray, corners: np.array):
		# convert from (n, 2) to (n, 1, 2)
		corners_original = []
		for corner in corners:
			corners_original.append([corner])

		corners_original = np.array(corners_original, dtype=np.float32)

		# Set the needed parameters to find the refined corners
		win_size = (5, 5)
		zero_zone = (-1, -1)
		criteria = (cv2.TERM_CRITERIA_EPS + cv2.TermCriteria_COUNT, 40, 0.001)

		# Calculate the refined corner locations
		corners_refined = cv2.cornerSubPix(src_gray, corners_original.copy(), win_size, zero_zone, criteria)

		# Write them down
		corner_result = []
		for _i in range(corners_original.shape[0]):
			corner = (corners_refined[_i, 0, 0], corners_refined[_i, 0, 1])
			corner_result.append(corner)

		return np.array(corner_result)


def project_load(path: str = "./"):
	_settings = None
	try:
		with open(path + "ocv_tracker_prj.json", "r") as fp:
			_settings = json.load(fp)
	except Exception:
		pass

	return _settings


def project_save(settings: dict, path: str = "./"):
	with open(path + "ocv_tracker_prj.json", "w") as fp:
		json.dump(settings, fp, indent=4)


if __name__ == '__main__':
	# Parse command line args
	parser = argparse.ArgumentParser(description="Gearbox Tracker")
	parser.add_argument("filename")
	parser.add_argument("--scale", default=1.0)
	parser.add_argument("--track", action="store_true")
	parser.add_argument("--path", action="store_true")
	parser.add_argument("--loop", action="store_true")
	parser.add_argument("--csv", action="store_true")
	parser.add_argument("--plot", action="store_true")
	parser.add_argument("--name", default="Default")
	args = parser.parse_args()
	video = cv2.VideoCapture(args.filename)

	# Fill CornerTrackerParams from  args
	ct_params = CornerTrackerParams(_scale=args.scale, _pre_track=args.track, _show_path=args.path)

	prj = project_load()
	prj_loaded = False
	if prj is not None:
		ct_params = ct_params.from_dict(prj['params'])
		prj_loaded = True

	# Let's go
	if not video.isOpened():
		print('Cannot open video file')
		sys.exit(1)
	# Read first frame.
	ok, image = video.read()
	if not ok:
		print('Cannot read video file')
		sys.exit(1)

	colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255)]
	tracker_count = 0
	tracker_list = []
	select_window = image.copy()

	if prj_loaded:
		ct_params = CornerTrackerParams().from_dict(prj['params'])
		for tracker_settings in prj['trackers']:
			print(tracker_settings)
			ct = CornerTracker(ct_params, colors[tracker_count], name=f"Tracker-{tracker_count}")
			ct.create_from_settings(select_window, CornerTrackerSettings().from_dict(tracker_settings))
			tracker_list.append(ct)
			tracker_count += 1

	else:
		prj = {'params': ct_params.__dict__, 'trackers': []}

	print("Press 'a' for adding trackers, press 'd' to delete a tracker, press other key to continue")
	k = cv2.waitKey(-1) & 0xff
	while True:
		print(f"Add tracker #{tracker_count}")
		ct = CornerTracker(ct_params, colors[tracker_count], name=f"Tracker-{tracker_count}")
		settings = ct.init_reference_frame(select_window)
		if settings is not None:
			prj['trackers'].append(settings.__dict__)
			tracker_list.append(ct)
			print(f"Number of active tracker: {len(tracker_list)}")
		else:
			break
		tracker_count += 1

	project_save(prj)

	# rewind to frame 0
	video.set(cv2.CAP_PROP_POS_FRAMES, 0)

	key_wait = -1
	dist_min = [(0, 0)]*len(tracker_list)
	dist_max = [(0, 0)]*len(tracker_list)
	result_csv = [[] for i in range(len(tracker_list))]
	result_plot = [[] for i in range(len(tracker_list))]
	while True:
		# Start timer
		timer = cv2.getTickCount()
		# Read a new frame
		ok, image = video.read()
		if not ok:
			if args.loop:
				video.set(cv2.CAP_PROP_POS_FRAMES, 0)
				continue
			else:
				break

		image_anno = image.copy()
		i = 0
		cv2.rectangle(image_anno, (25, 0), (int(image_anno.shape[1]), 25*(1+len(tracker_list))), (0, 0, 0), -1)
		for ct in tracker_list:
			mean_distance = ct.process(image, image_anno)
			if mean_distance is not None:
				scale = np.float32(ct_params.scale)
				scaled_distance = (scale*mean_distance[0], scale*mean_distance[1])
				if args.csv:
					result_csv[i].append({f'dx': scaled_distance[0], 'dy': scaled_distance[1]})
				if args.plot:
					result_plot[i].append(scaled_distance)
				dist_min[i] = (min(scaled_distance[0], dist_min[i][0]), min(scaled_distance[1], dist_min[i][1]))
				dist_max[i] = (max(scaled_distance[0], dist_max[i][0]), max(scaled_distance[1], dist_max[i][1]))
				cv2.putText(image_anno, f"Distance [{i}] : ({scaled_distance[0]:+05.2f}, "
							f"{scaled_distance[1]:+05.2f}), Min: ({dist_min[i][0]:+05.2f}, {dist_min[i][1]:+05.2f}),"
							f" Max: ({dist_max[i][0]:+05.2f}, {dist_max[i][1]:+05.2f})",
						(25, 25 * (i + 1)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, ct.color, 1)
			i += 1

		# Calculate Frames per second (FPS)
		fps = cv2.getTickFrequency() / (cv2.getTickCount() - timer)
		# Display FPS on frame
		cv2.putText(image_anno, "FPS : " + str(int(fps)), (20, image_anno.shape[0] - 20),
					cv2.FONT_HERSHEY_SIMPLEX, 0.5, (50, 170, 50), 1)

		cv2.imshow(f"Image Anno", image_anno)

		# Exit if ESC pressed
		k = cv2.waitKey(key_wait) & 0xff
		if k == 27:
			break
		if k == ord(' '):
			if key_wait == -1:
				key_wait = 1
			else:
				key_wait = -1

	# writing to csv file
	if len(result_csv[0]) > 0:
		i = 1
		for result in result_csv:
			filename = f"{args.name}_{i}.csv"
			with open(filename, 'w') as csvfile:
				# creating a csv dict writer object
				writer = csv.DictWriter(csvfile, fieldnames=['dx', 'dy'])

				# writing headers (field names)
				writer.writeheader()

				# writing data rows
				writer.writerows(result)
				print(f"Exported {filename}")
				i += 1

	if len(result_plot[0]) > 0:
		N_p = len(result_plot)
		i = 1
		for result in result_plot:
			Np = len
			n = range(0, len(result))
			plt.subplot(N_p, 1, i)
			plt.plot(n, np.array(result)[:, 0], n, np.array(result)[:, 1])
			plt.grid()
			plt.legend([f"x", f"y"])
			i += 1

	plt.show()