""" """

"""
This program is free software: you can redistribute it and/or modify it under
the terms of the GNU General Public License as published by the Free Software
Foundation, either version 3 of the License, or (at your option) any later
version.
This program is distributed in the hope that it will be useful, but WITHOUT
ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
You should have received a copy of the GNU General Public License along with
this program. If not, see <https://www.gnu.org/licenses/>.
"""

__author__ = "Théo de la Hogue"
__credits__ = []
__copyright__ = "Copyright 2023, Ecole Nationale de l'Aviation Civile (ENAC)"
__license__ = "GPLv3"

import unittest
import os
import math

from argaze.ArUcoMarker import ArUcoMarkerGroup, ArUcoMarker

import cv2 as cv
import numpy

class TestArUcoMarkerGroupClass(unittest.TestCase):

    def new_from_obj(self):

        # Edit file path
        current_directory = os.path.dirname(os.path.abspath(__file__))
        obj_filepath = os.path.join(current_directory, 'utils/scene.obj')

        # Load file
        self.aruco_markers_group = ArUcoMarkerGroup.ArUcoMarkerGroup.from_obj(obj_filepath)

    def new_from_json(self):

        # Edit file path
        current_directory = os.path.dirname(os.path.abspath(__file__))
        json_filepath = os.path.join(current_directory, 'utils/scene.json')

        # Load file
        self.aruco_markers_group = ArUcoMarkerGroup.ArUcoMarkerGroup.from_json(json_filepath)

    def setup_markers(self):

        # Prepare detected markers
        self.detected_markers = {
            0: ArUcoMarker.ArUcoMarker('DICT_ARUCO_ORIGINAL', 0, 1.),
            1: ArUcoMarker.ArUcoMarker('DICT_ARUCO_ORIGINAL', 1, 1.),
            2: ArUcoMarker.ArUcoMarker('DICT_ARUCO_ORIGINAL', 2, 1.),
            3: ArUcoMarker.ArUcoMarker('DICT_ARUCO_ORIGINAL', 3, 1.)
        }

        # Prepare scene markers and remaining markers
        self.scene_markers, self.remaining_markers = self.aruco_markers_group.filter_markers(self.detected_markers())

    def test_new_from_obj(self):
        """Test ArUcoMarkerGroup creation."""

        self.new_from_obj()
        self.setup_markers()

         # Check ArUcoMarkerGroup creation
        self.assertEqual(len(self.aruco_markers_group.places), 3)
        self.assertIsNone(numpy.testing.assert_array_equal(self.aruco_markers_group.identifiers, [0, 1, 2]))
        self.assertEqual(self.aruco_markers_group.marker_size, 1.)

        self.assertEqual(self.aruco_markers_group.places[0].marker.identifier, 0)
        self.assertIsNone(numpy.testing.assert_array_equal(self.aruco_markers_group.places[0].translation, [0., 0., 0.]))
        self.assertIsNone(numpy.testing.assert_array_equal(self.aruco_markers_group.places[0].rotation, [[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]))

        self.assertEqual(self.aruco_markers_group.places[1].marker.identifier, 1)
        self.assertIsNone(numpy.testing.assert_array_equal(self.aruco_markers_group.places[1].translation, [10., 10., 0.]))
        self.assertIsNone(numpy.testing.assert_array_equal(self.aruco_markers_group.places[1].rotation, [[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]))

        self.assertEqual(self.aruco_markers_group.places[2].marker.identifier, 2)
        self.assertIsNone(numpy.testing.assert_array_equal(self.aruco_markers_group.places[2].translation, [0., 10., 0.]))
        self.assertIsNone(numpy.testing.assert_array_equal(self.aruco_markers_group.places[2].rotation, [[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]))

    def test_new_from_json(self):
        """Test ArUcoMarkerGroup creation."""

        self.new_from_json()
        self.setup_markers()

         # Check ArUcoMarkerGroup creation
        self.assertEqual(len(self.aruco_markers_group.places), 3)
        self.assertIsNone(numpy.testing.assert_array_equal(self.aruco_markers_group.identifiers, [0, 1, 2]))
        self.assertEqual(self.aruco_markers_group.marker_size, 1.)

        self.assertEqual(self.aruco_markers_group.places[0].marker.identifier, 0)
        self.assertIsNone(numpy.testing.assert_array_equal(self.aruco_markers_group.places[0].translation, [0., 0., 0.]))
        self.assertIsNone(numpy.testing.assert_array_equal(self.aruco_markers_group.places[0].rotation, [[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]))

        self.assertEqual(self.aruco_markers_group.places[1].marker.identifier, 1)
        self.assertIsNone(numpy.testing.assert_array_equal(self.aruco_markers_group.places[1].translation, [10., 10., 0.]))
        self.assertIsNone(numpy.testing.assert_array_equal(self.aruco_markers_group.places[1].rotation, [[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]))

        self.assertEqual(self.aruco_markers_group.places[2].marker.identifier, 2)
        self.assertIsNone(numpy.testing.assert_array_equal(self.aruco_markers_group.places[2].translation, [0., 10., 0.]))
        self.assertIsNone(numpy.testing.assert_array_equal(self.aruco_markers_group.places[2].rotation, [[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]))

    def test_filter_markers(self):
        """Test ArUcoMarkerGroup markers filtering."""

        self.new_from_obj()
        self.setup_markers()

        # Check scene markers and remaining markers
        self.assertEqual(len(self.scene_markers), 3)
        self.assertEqual(len(self.remaining_markers), 1)

        self.assertIsNone(numpy.testing.assert_array_equal(list(self.scene_markers.keys()), self.aruco_markers_group.identifiers))
        self.assertIsNone(numpy.testing.assert_array_equal(list(self.remaining_markers.keys()), [3]))

    def test_check_markers_consistency(self):
        """Test ArUcoMarkerGroup markers consistency checking."""

        self.new_from_obj()
        self.setup_markers()

        # Edit consistent marker poses
        self.scene_markers[0].translation = numpy.array([1., 1., 5.])
        self.scene_markers[0].rotation = numpy.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])

        self.scene_markers[1].translation = numpy.array([11., 11., 5.])
        self.scene_markers[1].rotation = numpy.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])

        self.scene_markers[2].translation = numpy.array([1., 11., 5.])
        self.scene_markers[2].rotation = numpy.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])

        # Check consistency
        consistent_markers, unconsistent_markers, unconsistencies = self.aruco_markers_group.check_markers_consistency(self.scene_markers, 1, 1)

        # Check consistent markers, unconsistent markers and unconsistencies
        self.assertEqual(len(consistent_markers), 3)
        self.assertEqual(len(unconsistent_markers), 0)
        self.assertEqual(len(unconsistencies['rotation']), 0)
        self.assertEqual(len(unconsistencies['translation']), 0)

        self.assertIsNone(numpy.testing.assert_array_equal(list(consistent_markers.keys()), self.aruco_markers_group.identifiers))

        # Edit unconsistent marker poses
        self.scene_markers[2].translation = numpy.array([5., 15., 5.])

        # Check consistency
        consistent_markers, unconsistent_markers, unconsistencies = self.aruco_markers_group.check_markers_consistency(self.scene_markers, 1, 1)

        # Check consistent markers, unconsistent markers and unconsistencies
        self.assertEqual(len(consistent_markers), 2)
        self.assertEqual(len(unconsistent_markers), 1)
        self.assertEqual(len(unconsistencies['rotation']), 0)
        self.assertEqual(len(unconsistencies['translation']), 2)

        self.assertIsNone(numpy.testing.assert_array_equal(list(unconsistent_markers.keys()), [2]))
        self.assertIsNone(numpy.testing.assert_array_equal(list(unconsistencies['translation'].keys()), ['0/2', '1/2']))
        self.assertIsNone(numpy.testing.assert_array_equal(list(unconsistencies['translation']['0/2'].keys()), ['current', 'expected']))
        self.assertIsNone(numpy.testing.assert_array_equal(list(unconsistencies['translation']['1/2'].keys()), ['current', 'expected']))

    def test_estimate_pose_from_single_marker(self):
        """Test ArUcoMarkerGroup pose estimation from single marker."""

        self.new_from_obj()
        self.setup_markers()

        # Edit marke pose
        self.scene_markers[0].translation = numpy.array([1., 1., 5.])
        self.scene_markers[0].rotation = numpy.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])

        # Estimate pose
        tvec, rmat = self.aruco_markers_group.estimate_pose_from_single_marker(self.scene_markers[0])

        self.assertIsNone(numpy.testing.assert_array_equal(tvec, [1., 1., 5.]))
        self.assertIsNone(numpy.testing.assert_array_equal(rmat, [[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]))

    def test_estimate_pose_from_markers(self):
        """Test ArUcoMarkerGroup pose estimation from markers."""

        self.new_from_obj()
        self.setup_markers()

        # Edit markers pose
        self.scene_markers[0].translation = numpy.array([1., 1., 5.])
        self.scene_markers[0].rotation = numpy.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])

        self.scene_markers[1].translation = numpy.array([11., 11., 5.])
        self.scene_markers[1].rotation = numpy.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])

        self.scene_markers[2].translation = numpy.array([1., 11., 5.])
        self.scene_markers[2].rotation = numpy.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])

        # Estimate pose
        tvec, rmat = self.aruco_markers_group.estimate_pose_from_markers(self.scene_markers)

        self.assertIsNone(numpy.testing.assert_array_equal(tvec, [1., 1., 5.]))
        self.assertIsNone(numpy.testing.assert_array_equal(rmat, [[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]))

    @unittest.skip("ArUcoMarkerGroup estimate_pose_from_axis_markers method is broken.")
    def test_estimate_pose_from_axis_markers(self):
        """Test ArUcoMarkerGroup pose estimation from axis markers."""

        self.new_from_obj()
        self.setup_markers()

        # Edit markers pose
        self.scene_markers[0].translation = numpy.array([1., 1., 5.])
        self.scene_markers[0].rotation = numpy.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])

        self.scene_markers[1].translation = numpy.array([11., 11., 5.])
        self.scene_markers[1].rotation = numpy.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])

        self.scene_markers[2].translation = numpy.array([1., 11., 5.])
        self.scene_markers[2].rotation = numpy.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])

        # Estimate pose
        tvec, rmat = self.aruco_markers_group.estimate_pose_from_axis_markers(self.scene_markers[2], self.scene_markers[1], self.scene_markers[0])

        self.assertIsNone(numpy.testing.assert_array_equal(tvec, [1., 1., 5.]))
        self.assertIsNone(numpy.testing.assert_array_equal(rmat, [[1., 0., 0.], [0., -1., 0.], [0., 0., -1.]]))

if __name__ == '__main__':

    unittest.main()