import sys
import numpy as np
import cv2 as cv
import math
import argparse
from utils import to_board_pose_euler, cap_init
from aruco_types import ARUCO_DICT

def factory(dict_type: int, square_length: float=5.0, marker_length: float=3.0, margin_length=0):
	aruco_dict = cv.aruco.getPredefinedDictionary(dict_type)
	board = cv.aruco.CharucoBoard((3,3), square_length, marker_length, aruco_dict)
	params = cv.aruco.CharucoParameters()
	detector = cv.aruco.CharucoDetector(board, params)
	image = board.generateImage((800, 600), None, margin_length, 1)
	cv.imwrite("results/charuco_diamond_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.CharucoDetector, board: cv.aruco.CharucoBoard):
	with np.load('results/cam_calib_charuco.npz') as X:
		mtx, dist, _, _ = [X[i] for i in ('mtx', 'dist', 'rvecs', 'tvecs')]

	# Open camera
	cap = cv.VideoCapture(0)
	if not cap.isOpened():
		print("Cannot open camera")
		exit()

	cap_init(cap)

	last_rotation = [0,0,0]
	while True:
		# Read a new frame
		ok, img = cap.read()
		if not ok:
			return
		k = cv.waitKey(1) & 0xFF
		if k == ord('q'):
			break
		else:
			r_vec, t_vec = process(detector, board, img, mtx, dist, None, None)
			if r_vec is not None and t_vec is not None:
				r_vec_obj, _ = to_board_pose_euler(r_vec, t_vec)
				rot = [round(180.0 / math.pi * v, 0) for v in r_vec_obj]
				if rot != last_rotation:
					print("\r                                                                                 ", end='')
					print(f"\rTilt: {rot[0]:.1f}, Roll: {rot[1]:.1f}, Azimuth: {rot[2]:.1f}", end='')
				last_rotation = rot

		cv.imshow('img', img)

def process(detector: cv.aruco.CharucoDetector, board: cv.aruco.CharucoBoard, img, mtx, dist, rvecs=None, tvecs=None, draw_marker_box=True, draw_marker_id=False, draw_marker_axis=False):	# Load previously saved data
	use_detect_diamonds = 0
	r_vec = None
	t_vec = None
	gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
	diamond_corners, diamond_ids, marker_corners, marker_ids = detector.detectBoard(gray)
	if diamond_corners is not None:
		if use_detect_diamonds:
			diamond_corners1, diamond_ids1, marker_corners1, marker_ids1 = detector.detectDiamonds(gray)
			diamond_corners = diamond_corners1[0]
			diamond_ids = np.transpose(diamond_ids1[0])
			marker_corners = marker_corners1
			marker_ids = marker_ids1[0]

		# Draw marker box
		if draw_marker_box:
			cv.aruco.drawDetectedMarkers(img, marker_corners)

		# Draw marker id
		if draw_marker_id:
			cv.aruco.drawDetectedDiamonds(img, diamond_corners, diamond_ids)

		# Draw marker axis
		if draw_marker_axis:
			_rvecs, _tvecs, _ = cv.aruco.estimatePoseSingleMarkers(marker_corners, markerLength=0.1, cameraMatrix=mtx, distCoeffs=dist)
			for i in range(len(_rvecs)):
				cv.drawFrameAxes(img, mtx, dist, _rvecs[i], _tvecs[i], 0.05)

		if len(diamond_ids) >= len(board.getIds()):
			# Estimate diamond pose
			obj_points, img_points = board.matchImagePoints(diamond_corners, diamond_ids)
			offset = int(board.getMarkerLength() * board.getSquareLength())
			obj_points -= [offset/2, offset/2, 0]
			if mtx is not None and dist is not None:
				pose, r_vec, t_vec = cv.solvePnP(obj_points, img_points, mtx, dist, useExtrinsicGuess=False, rvec=rvecs,
												 tvec=tvecs, flags=cv.SOLVEPNP_ITERATIVE)
				if pose:
					# Draw Axis
					cv.drawFrameAxes(img, mtx, dist, r_vec, t_vec, 10)

	return r_vec, t_vec

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