import argparse
import cv2
import sys

import cv2.aruco

from aruco_types import ARUCO_DICT


def im_resize(image, new_width: int):
	# Define new width while maintaining the aspect ratio
	aspect_ratio = new_width / image.shape[0]
	new_height = int(image.shape[1] * aspect_ratio)  # Compute height based on aspect ratio

	image = cv2.resize(image, (new_width, new_height))
	return image

def factory(aruco_type: int) -> cv2.aruco.ArucoDetector:
	aruco_dict = cv2.aruco.getPredefinedDictionary(aruco_type)
	aruco_params = cv2.aruco.DetectorParameters()
	aruco_detector = cv2.aruco.ArucoDetector(aruco_dict, aruco_params)

	return aruco_detector

def process_video(aruco_detector: cv2.aruco.ArucoDetector):
	# Open camera
	video = cv2.VideoCapture(0)
	if not video.isOpened():
		print("Cannot open camera")
		exit()

	while True:
		# Read a new frame
		ok, img = video.read()
		if not ok:
			return
		process(aruco_detector, img)
		k = cv2.waitKey(1) & 0xFF
		if k == ord('q'):
			break

def process_still(aruco_detector, filename: str):
	image = cv2.imread(filename)
	image = im_resize(image, 600)
	process(aruco_detector, image)
	k = cv2.waitKey(-1) & 0xFF

def process(aruco_detector, img):
	(corners, ids, rejected) = aruco_detector.detectMarkers(img)
	img = cv2.aruco.drawDetectedMarkers(img, corners, ids)
	cv2.imshow("Image", img)

def main():
	# construct the argument parser and parse the arguments
	ap = argparse.ArgumentParser()
	ap.add_argument("-i", "--image",
		help="path to input image containing ArUCo tag")
	ap.add_argument("-t", "--type", type=str,
		default="DICT_ARUCO_ORIGINAL",
		help="type of ArUCo tag to detect")
	args = vars(ap.parse_args())

	# verify that the supplied ArUCo tag exists and is supported by
	# OpenCV
	if ARUCO_DICT.get(args["type"], None) is None:
		print("[INFO] ArUCo tag of '{}' is not supported".format(
			args["type"]))
		sys.exit(0)
	# load the ArUCo dictionary, grab the ArUCo parameters, and detect
	# the markers
	print("[INFO] detecting '{}' tags...".format(args["type"]))

	aruco_detector = factory(ARUCO_DICT[args["type"]])

	# load the input image from disk and resize it
	print("[INFO] loading image...")

	if args["image"] is not None:
		process_still(aruco_detector, args["image"])
	else:
		process_video(aruco_detector)


if __name__ == '__main__':
	main()