import cv2
import numpy as np

CLIPPING_MASK_MIN = 0.3
CLIPPING_MASK_MAX = 2.0
CLIPPING_MASK_NEAR = 0.9
CLIPPING_MASK_FAR = 1.0

EPSILON = 1e-20
ARROWLENGTH = 20

# State variables
_prev_pos = None
_prev_heading = None

def obtain_depth_view(image, size):
	frame = image.copy()
	frame = cv2.resize(frame, (size, size))
	skewmin = np.amin(frame)
	skewmax = np.amax(frame)

	np.clip(frame, CLIPPING_MASK_MIN, CLIPPING_MASK_MAX, out=frame)
	frame -= skewmin
	frame /= (skewmax - skewmin) / 255.0
	frame = frame.astype(dtype=np.uint8, copy=False)

	im_color = cv2.applyColorMap(frame, cv2.COLORMAP_BONE)
	return im_color


def process(image, size):
	repr_image = obtain_depth_view(image, size)

	depth_image_3d = image.copy()
	depth_image_3d = cv2.resize(depth_image_3d, (size, size))

	white = 255 * np.ones(depth_image_3d.shape, np.uint8)
	bg_removed = np.where((depth_image_3d > CLIPPING_MASK_FAR) | (depth_image_3d < CLIPPING_MASK_NEAR), 0, white)

	blurred = cv2.blur(bg_removed, (5, 5))
	kernel = np.ones((4, 4), np.uint8)
	dilation = cv2.dilate(blurred, kernel, iterations=5)
	kernel = np.ones((5, 5), np.uint8)
	erosion = cv2.erode(dilation, kernel, iterations=3)
	# edged = cv2.Canny(erosion, 100, 200, 5)

	contours, _ = cv2.findContours(erosion, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
	csize = len(contours)
	if csize == 0:
		print("No Contours Detected!")
		return repr_image, (0, 0, 0)
	else:
		for c in contours:
			x, y, w, h = cv2.boundingRect(c)
			cv2.rectangle(repr_image, (x, y), (x + w, y + h), (255, 0, 90), 2)

		ctr, ccenter, tl, br = _findLargestContour(contours)
		cv2.circle(repr_image, ccenter, 4, (255, 0, 0), -1)

		heading = 0

		global _prev_pos, _prev_heading
		if _prev_pos is not None:
			# You have moved if the distance is large enough
			#   small distances can occur due to the object detection
			pdist = _tuple_distance(ccenter, _prev_pos)
			if pdist < EPSILON:
				# You have not actually moved
				ccenter = _prev_pos
			else:
				# pv is the normalized direction identified in the robot
				pv = (ccenter[0] - _prev_pos[0]) / pdist, (ccenter[1] - _prev_pos[1]) / pdist
				heading = np.arctan2(pv[1], pv[0])

				cv2.arrowedLine(repr_image, ccenter, (ccenter[0] + int(pv[0] * ARROWLENGTH), ccenter[1] + int(pv[1] * ARROWLENGTH)),
				                (255, 0, 0), 2, tipLength=0.3)

				if _prev_heading is not None:
					# inconsistencies in the object detection may result in an opposite angle
					angle_diff = abs(_prev_heading - heading)
					while angle_diff > np.pi / 2:
						if heading > _prev_heading:
							heading -= np.pi
						else:
							heading += np.pi
						angle_diff = abs(_prev_heading - heading)

		_prev_heading = heading
		_prev_pos = ccenter

	# field of view transformation
	#   NOTE: the obtained image uses a cone for the distance, it is the
	#         vertical component instead!
	mid = int(size / 2), int(size / 2)
	cd = image[mid[0]][mid[1]][0]
	try:
		pd = image[ccenter[0]][ccenter[1]][0]
		mid_xy = (mid[0] - ccenter[0]) ** 2 + (mid[1] + ccenter[1]) ** 2
		pd = np.sqrt(mid_xy - pd ** 2)
	except:
		pd = 0
	if pd == 0 or pd > cd:
		pd = (CLIPPING_MASK_NEAR + CLIPPING_MASK_FAR) / 2

	ratio = 1
	if cd != 0:
		ratio = pd / cd
	ncenter = (np.asarray(ccenter) - np.asarray(mid)) * ratio + np.asarray(mid)

	cv2.circle(repr_image, (int(ncenter[0]), int(ncenter[1])), 4, (255, 100, 0), -1)

	# NOTE: the origin is the top left corner of the image!
	return repr_image, (ncenter[0] / size, ncenter[1] / size, heading)


def _findLargestContour(contours):
	cnt = sorted(contours, key=cv2.contourArea)[-1]
	M = cv2.moments(cnt)
	pos = int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"])
	x, y, w, h = cv2.boundingRect(cnt)
	return cnt, pos, (x, y), (x + w, y + h)


def _tuple_distance(p1, p2):
	return np.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)