#!/usr/bin/env python

import cv2 as cv
import cv2.aruco as aruco
import numpy

all_aruco_markers_dictionaries = {
    'DICT_4X4_50': aruco.DICT_4X4_50,
    'DICT_4X4_100': aruco.DICT_4X4_100,
    'DICT_4X4_250': aruco.DICT_4X4_250,
    'DICT_4X4_1000': aruco.DICT_4X4_1000,
    'DICT_5X5_50': aruco.DICT_5X5_50,
    'DICT_5X5_100': aruco.DICT_5X5_100,
    'DICT_5X5_250': aruco.DICT_5X5_250,
    'DICT_5X5_1000': aruco.DICT_5X5_1000,
    'DICT_6X6_50': aruco.DICT_6X6_50,
    'DICT_6X6_100': aruco.DICT_6X6_100,
    'DICT_6X6_250': aruco.DICT_6X6_250,
    'DICT_6X6_1000': aruco.DICT_6X6_1000,
    'DICT_7X7_50': aruco.DICT_7X7_50,
    'DICT_7X7_100': aruco.DICT_7X7_100,
    'DICT_7X7_250': aruco.DICT_7X7_250,
    'DICT_7X7_1000': aruco.DICT_7X7_1000,
    'DICT_ARUCO_ORIGINAL': aruco.DICT_ARUCO_ORIGINAL,
    'DICT_APRILTAG_16h5': aruco.DICT_APRILTAG_16h5,
    'DICT_APRILTAG_25h9': aruco.DICT_APRILTAG_16h5,
    'DICT_APRILTAG_36h10': aruco.DICT_APRILTAG_36h10,
    'DICT_APRILTAG_36h11': aruco.DICT_APRILTAG_36h11
}
"""Dictionnary to list all built-in ArUco markers dictionaries from OpenCV package."""

class ArUcoMarkersDictionary():
    """Handle an ArUco markers dictionary."""

    def __init__(self, aruco_dictionary_name):

        if all_aruco_markers_dictionaries.get(aruco_dictionary_name, None) is None:
            raise NameError(f'Bad ArUco markers dictionary name: {aruco_dictionary_name}')

        dict_name_split = aruco_dictionary_name.split('_')

        self.__format = dict_name_split[1]

        # DICT_ARUCO_ORIGINAL case
        if self.__format == 'ARUCO':

            self.__format = '5X5'
            self.__number = 1024

        # DICT_APRILTAG_ case
        elif self.__format == 'APRILTAG':

            april_tag_format = dict_name_split[2]

            if april_tag_format == '16h5':

                self.__format = '4X4'
                self.__number = 30

            elif april_tag_format == '25h9':

                self.__format = '5X5'
                self.__number = 30

            elif april_tag_format == '36h10':

                self.__format = '6X6'
                self.__number = 2320

            elif april_tag_format == '36h11':

                self.__format = '6X6'
                self.__number = 587
  
        # other cases
        else:

            self.__number = int(dict_name_split[2])

        self.__aruco_dict = aruco.Dictionary_get(all_aruco_markers_dictionaries[aruco_dictionary_name])

    def create_marker(self, i, dpi=300):
        """Create a marker image."""

        marker = numpy.zeros((dpi, dpi, 1), dtype="uint8")
        aruco.drawMarker(self.__aruco_dict, i, dpi, marker, 1)

        return numpy.repeat(marker, 3).reshape(dpi, dpi, 3)

    def get_markers(self):
        return self.__aruco_dict

    def get_markers_format(self):
        return self.__format

    def get_markers_number(self):
        return self.__number

    def export_as_png(self, destination_folder, dpi, i):
        """Save one marker into a .png file."""

        if i >= 0 and i < self.__number:

            output_filename = f'marker_{self.__format}_{i}.png'

            print(destination_folder, output_filename)

            # create marker
            marker = self.create_marker(i, dpi)

            # save marker into destination folder
            cv.imwrite(f'{destination_folder}/{output_filename}', marker)

        else:
            raise ValueError(f'Bad ArUco index: {i}')

    def export_all(self, destination_folder, dpi):
        """Save all markers dictionary into separated .png files."""

        for i in range(self.__number):

            self.export_as_png(destination_folder, dpi, i)