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

def generate(aruco_dict: cv2.aruco.Dictionary, tag_id: int, side_pixels: int = 300):
	tag = np.zeros((side_pixels, side_pixels, 1), dtype="uint8")
	cv2.aruco.generateImageMarker(aruco_dict, tag_id, side_pixels, tag, 1)

	return tag

def main():
	# construct the argument parser and parse the arguments
	ap = argparse.ArgumentParser()
	ap.add_argument("-f", "--folder", required=True,
					help="path to output image containing ArUCo tag")
	ap.add_argument("-d", "--id", type=int, required=True,
					help="ID of ArUCo tag to generate")
	ap.add_argument("-t", "--type", type=str,
					default="DICT_ARUCO_ORIGINAL",
					help="type of ArUCo tag to generate")
	ap.add_argument("-i", "--imagetype", type=str,
					default="png",
					help="image file type of ArUCo tag image")
	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
	aruco_dict = cv2.aruco.getPredefinedDictionary(ARUCO_DICT[args["type"]])

	# allocate memory for the output ArUCo tag and then draw the ArUCo
	# tag on the output image
	print("[INFO] generating ArUCo tag type '{}' with ID '{}'".format(
		args["type"], args["id"]))
	tag = generate(aruco_dict, args["id"], 300)
	# write the generated ArUCo tag to disk and then display it to our
	# screen
	filename = f"{args["folder"]}/{args["type"]}_ID_{args["id"]:04d}.{args["imagetype"]}"
	cv2.imwrite(filename, tag)
	cv2.imshow("ArUCo Tag", tag)
	cv2.waitKey(0)

if __name__ == '__main__':
	main()