import sys
import numpy as np
import cv2 as cv
import argparse
from aruco_types import ARUCO_DICT

def factory(aruco_type: int):
	aruco_dict = cv.aruco.getPredefinedDictionary(aruco_type)
	params = cv.aruco.DetectorParameters()
	detector = cv.aruco.ArucoDetector(aruco_dict, params)

	board = cv.aruco.GridBoard(size=(5, 7), markerLength=4, markerSeparation=2, dictionary=aruco_dict)
	image = np.zeros((800,600), np.uint8)
#	image = board.generateImage((800, 600), None, 0, 1)
	cv.imwrite("results/aruco_board.png", image)
	return detector, board

def main():
	# construct the argument parser and parse the arguments
	ap = argparse.ArgumentParser()
	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"]))
	detector, board = factory(ARUCO_DICT[args["type"]])

	# load the input image from disk and resize it
	print("[INFO] start processing...")
	process_video(detector, board)

def process_video(detector: cv.aruco.ArucoDetector, board):
	with np.load('results/cam.npz') as X:
		mtx, dist, _, _ = [X[i] for i in ('mtx', 'dist', 'rvecs', 'tvecs')]

	# Open camera
	video = cv.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(detector, board, img, mtx, dist)
		k = cv.waitKey(1) & 0xFF
		if k == ord('q'):
			break

def process(detector: cv.aruco.ArucoDetector, board: cv.aruco.Board, img, mtx, dist):	# Load previously saved data
	gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
	(corners, ids, rejected) = detector.detectMarkers(gray)

	detector.refineDetectedMarkers(img, board, corners, ids, rejected, mtx, dist)

	# Draw a square around the markers
	cv.aruco.drawDetectedMarkers(img, corners)

	if len(corners) > 0:
		# Estimate pose of each marker and return the values rvec and tvec---(different from those of camera coefficients)
		obj_points, img_points = board.matchImagePoints(corners, ids)
		pose, rvec, tvec = cv.solvePnP(obj_points, img_points, mtx, dist, None, None)
		if pose:
			# Draw Axis
			cv.drawFrameAxes(img, mtx, dist, rvec, tvec, 20)


	cv.imshow('img', img)


if __name__ == '__main__':
	main()
	cv.destroyAllWindows()