#!/usr/bin/env python

""" """

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

from typing import TypeVar, Tuple
from dataclasses import dataclass, field
import json
import math
import itertools
import re

from argaze import DataFeatures
from argaze.ArUcoMarkers import ArUcoMarkersDictionary, ArUcoMarker, ArUcoOpticCalibrator

import numpy
import cv2

T0 = numpy.array([0., 0., 0.])
"""Define no translation vector."""

R0 = numpy.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])
"""Define no rotation matrix."""

ArUcoMarkersGroupType = TypeVar('ArUcoMarkersGroup', bound="ArUcoMarkersGroup")
# Type definition for type annotation convenience

def make_rotation_matrix(x, y, z):

    # Create rotation matrix around x axis
    c = numpy.cos(numpy.deg2rad(x))
    s = numpy.sin(numpy.deg2rad(x))
    Rx = numpy.array([[1, 0, 0], [0, c, -s], [0, s, c]])

    # Create rotation matrix around y axis
    c = numpy.cos(numpy.deg2rad(y))
    s = numpy.sin(numpy.deg2rad(y))
    Ry = numpy.array([[c, 0, s], [0, 1, 0], [-s, 0, c]])

    # Create rotation matrix around z axis
    c = numpy.cos(numpy.deg2rad(z))
    s = numpy.sin(numpy.deg2rad(z))
    Rz = numpy.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])

    # Return intrinsic rotation matrix
    return Rx.dot(Ry.dot(Rz))

def is_rotation_matrix(R):

    Rt = numpy.transpose(R)
    shouldBeIdentity = numpy.dot(Rt, R)
    I = numpy.identity(3, dtype = R.dtype)
    n = numpy.linalg.norm(I - shouldBeIdentity)

    return n < 1e-3

@dataclass(frozen=True)
class Place():
    """Define a place as list of corners position and a marker.

    Parameters:
        corners: 3D corners position in group referential.
        marker: ArUco marker linked to the place.
    """

    corners: numpy.array
    marker: dict

@dataclass
class ArUcoMarkersGroup(DataFeatures.PipelineStepObject):
    """
    Handle group of ArUco markers as one unique spatial entity and estimate its pose.
    """

    def __init__(self, dictionary: ArUcoMarkersDictionary.ArUcoMarkersDictionary = None, places: dict = None, **kwargs):
        """Initialize ArUcoMarkersGroup

        Parameters:
            dictionary: expected dictionary of all markers in the group.
            places: expected markers place.
        """

        # Init parent classes
        DataFeatures.PipelineStepObject.__init__(self, **kwargs)

        # Init private attributes
        self.__dictionary = dictionary
        self.__places = places
        self.__translation = numpy.zeros(3)
        self.__rotation = numpy.zeros(3)

        # Normalize places data
        new_places = {}

        for identifier, data in self.__places.items():

            # Convert string identifier to int value
            if type(identifier) == str:

                identifier = int(identifier)

                # Get translation vector
                tvec = numpy.array(data.pop('translation')).astype(numpy.float32)

                # Check rotation value shape
                rvalue = numpy.array(data.pop('rotation')).astype(numpy.float32)

                # Rotation matrix
                if rvalue.shape == (3, 3):

                    rmat = rvalue

                # Rotation vector (expected in degree)
                elif rvalue.shape == (3,):

                    rmat = make_rotation_matrix(rvalue[0], rvalue[1], rvalue[2]).astype(numpy.float32)

                else:

                    raise ValueError(f'Bad rotation value: {rvalue}')

                assert(is_rotation_matrix(rmat))

                # Get marker size
                size = numpy.array(data.pop('size')).astype(numpy.float32)

                new_marker = ArUcoMarker.ArUcoMarker(self.__dictionary, identifier, size)

                # Build marker corners thanks to translation vector and rotation matrix
                place_corners = numpy.array([[-size/2, size/2, 0], [size/2, size/2, 0], [size/2, -size/2, 0], [-size/2, -size/2, 0]])
                place_corners = place_corners.dot(rmat) + tvec

                new_places[identifier] = Place(place_corners, new_marker)

            # else places are configured using detected markers estimated points
            elif isinstance(data, ArUcoMarker.ArUcoMarker):

                new_places[identifier] = Place(data.points, data)

            # else places are already at expected format
            elif (type(identifier) == int) and isinstance(data, Place):

                new_places[identifier] = data

        self.__places = new_places

    @property
    def dictionary(self) -> ArUcoMarkersDictionary.ArUcoMarkersDictionary:
        """Get ArUco marker group ArUco dictionary."""
        return self.__dictionary
    
    @property
    def places(self) -> dict:
        """Get ArUco marker group places dictionary."""
        return self.__places

    @property
    def identifiers(self) -> list:
        """List place marker identifiers belonging to the group."""
        return list(self.__places.keys())

    @property
    def translation(self) -> numpy.array:
        """Get ArUco marker group translation vector."""
        return self.__translation

    @translation.setter
    def translation(self, tvec):
        """Set ArUco marker group translation vector."""
        self.__translation = tvec

    @property
    def rotation(self) -> numpy.array:
        """Get ArUco marker group rotation matrix."""
        return self.__translation

    @rotation.setter
    def rotation(self, rmat):
        """Set ArUco marker group rotation matrix."""
        self.__rotation = rmat

    def as_dict(self) -> dict:
        """Export ArUco marker group properties as dictionary."""

        return {
            **DataFeatures.PipelineStepObject.as_dict(self),
            "dictionary": self.__dictionary,
            "places": self.__places
        }
        
    @classmethod
    def from_dict(cls, aruco_markers_group_data: dict, working_directory: str = None) -> ArUcoMarkersGroupType:
        """Load ArUco markers group attributes from dictionary.

        Parameters:
            aruco_markers_group_data: dictionary with attributes to load
            working_directory: folder path where to load files when a dictionary value is a relative filepath.
        """

        new_dictionary = aruco_markers_group_data.pop('dictionary')
        new_places = aruco_markers_group_data.pop('places')

        return ArUcoMarkersGroup(new_dictionary, new_places)

    @classmethod
    def from_obj(self, obj_filepath: str) -> ArUcoMarkersGroupType:
        """Load ArUco markers group from .obj file.

        !!! note
            Expected object (o) name format: <DICTIONARY>#<IDENTIFIER>_Marker

        !!! note
            All markers have to belong to the same dictionary.

        """

        new_dictionary = None
        new_places = {}
        
        # Regex rules for .obj file parsing
        OBJ_RX_DICT = {
            'object': re.compile(r'o (.*)#([0-9]+)_(.*)\n'),
            'vertice': re.compile(r'v ([+-]?[0-9]*[.]?[0-9]+) ([+-]?[0-9]*[.]?[0-9]+) ([+-]?[0-9]*[.]?[0-9]+)\n'),
            'face': re.compile(r'f ([0-9]+) ([0-9]+) ([0-9]+) ([0-9]+)\n'),
            'comment': re.compile(r'#(.*)\n') # keep comment regex after object regex because the # is used in object string too
        }

        # Regex .obj line parser
        def __parse_obj_line(line):

            for key, rx in OBJ_RX_DICT.items():
                match = rx.search(line)
                if match:
                    return key, match

            # If there are no matches
            return None, None
        
        # Start parsing
        try:

            identifier = None
            vertices = []
            faces = {}

            # Open the file and read through it line by line
            with open(obj_filepath, 'r') as file:

                line = file.readline()

                while line:

                    # At each line check for a match with a regex
                    key, match = __parse_obj_line(line)

                    # Extract comment
                    if key == 'comment':
                        pass

                    # Extract marker dictionary and identifier
                    elif key == 'object':

                        dictionary = str(match.group(1))
                        identifier = int(match.group(2))
                        last = str(match.group(3))

                        # Init new group dictionary with first dictionary name
                        if new_dictionary == None:

                            new_dictionary = ArUcoMarkersDictionary.ArUcoMarkersDictionary(dictionary)

                        # Check all others marker dictionary are equal to new group dictionary 
                        elif dictionary != new_dictionary.name:

                            raise NameError(f'Marker {identifier} dictionary is not {new_dictionary.name}')

                    # Fill vertices array
                    elif key == 'vertice':

                        vertices.append(tuple([float(match.group(1)), float(match.group(2)), float(match.group(3))]))

                    # Extract vertice ids
                    elif key == 'face':

                        faces[identifier] = [int(match.group(1)), int(match.group(2)), int(match.group(3)), int(match.group(4))]

                    # Go to next line
                    line = file.readline()

                file.close()

                # Retreive marker vertices thanks to face vertice ids
                for identifier, face in faces.items():

                    # Gather place corners in clockwise order
                    cw_corners = numpy.array([ vertices[i-1] for i in reversed(face) ])

                    # Edit place axis from corners positions
                    place_x_axis = cw_corners[2] - cw_corners[3]
                    place_x_axis_norm = numpy.linalg.norm(place_x_axis)
  
                    place_y_axis = cw_corners[0] - cw_corners[3]
                    place_y_axis_norm = numpy.linalg.norm(place_y_axis)

                    # Check axis size: they should be almost equal
                    if math.isclose(place_x_axis_norm, place_y_axis_norm, rel_tol=1e-3):

                        new_marker_size = place_x_axis_norm

                    else:

                        raise ValueError(f'{new_dictionary}#{identifier}_Marker is not a square.')

                    # Create a new place related to a new marker
                    new_marker = ArUcoMarker.ArUcoMarker(new_dictionary, identifier, new_marker_size)
                    new_places[identifier] = Place(cw_corners, new_marker)

        except IOError:
            raise IOError(f'File not found: {obj_filepath}')

        return ArUcoMarkersGroup(new_dictionary, new_places)

    def filter_markers(self, detected_markers: dict) -> Tuple[dict, dict]:
        """Sort markers belonging to the group from given detected markers dict (cf ArUcoDetector.detect_markers()).

        Returns:
            dict of markers belonging to this group
            dict of remaining markers not belonging to this group
        """

        group_markers = {}
        remaining_markers = {}

        for (marker_id, marker) in detected_markers.items():

            if marker_id in self.__places.keys():

                group_markers[marker_id] = marker

            else:
                
                remaining_markers[marker_id] = marker

        return group_markers, remaining_markers

    def estimate_pose_from_markers_corners(self, markers: dict, K: numpy.array, D: numpy.array) -> Tuple[bool, numpy.array, numpy.array]:
        """Estimate pose from markers corners and places corners.

        Parameters:
            markers: detected markers to use for pose estimation.
            K: intrinsic camera parameters
            D: camera distorsion matrix

        Returns:
            success: True if the pose estimation succeeded
            tvec: scene translation vector
            rvec: scene rotation vector
        """

        markers_corners_2d = []
        places_corners_3d = []

        for identifier, marker in markers.items():

            try:

                place = self.__places[identifier]

                for marker_corner in marker.corners:
                    markers_corners_2d.append(list(marker_corner))

                for place_corner in place.corners:
                    places_corners_3d.append(list(place_corner))

            except KeyError:

                raise ValueError(f'Marker {marker.identifier} doesn\'t belong to the group.')

        # SolvPnP using cv2.SOLVEPNP_SQPNP flag
        # TODO: it works also with cv2.SOLVEPNP_EPNP flag so we need to test which is the faster.
        # About SolvPnP flags: https://docs.opencv.org/4.x/d5/d1f/calib3d_solvePnP.html
        success, rvec, tvec = cv2.solvePnP(numpy.array(places_corners_3d), numpy.array(markers_corners_2d), numpy.array(K), numpy.array(D), flags=cv2.SOLVEPNP_SQPNP)

        # Refine pose estimation using Gauss-Newton optimisation
        if success :

            rvec, tvec = cv2.solvePnPRefineVVS(numpy.array(places_corners_3d), numpy.array(markers_corners_2d), numpy.array(K), numpy.array(D), rvec, tvec)

        self.__translation = tvec.T
        self.__rotation = rvec.T

        return success, self.__translation, self.__rotation

    def draw_axes(self, image: numpy.array, K, D, thickness: int = 0, length: float = 0):
        """Draw group axes."""

        try:
            axisPoints = numpy.float32([[length, 0, 0], [0, length, 0], [0, 0, length], [0, 0, 0]]).reshape(-1, 3)
            axisPoints, _ = cv2.projectPoints(axisPoints, self.__rotation, self.__translation, numpy.array(K), numpy.array(D))
            axisPoints = axisPoints.astype(int)

            cv2.line(image, tuple(axisPoints[3].ravel()), tuple(axisPoints[0].ravel()), (0, 0, 255), thickness) # X (red)
            cv2.line(image, tuple(axisPoints[3].ravel()), tuple(axisPoints[1].ravel()), (0, 255, 0), thickness) # Y (green)
            cv2.line(image, tuple(axisPoints[3].ravel()), tuple(axisPoints[2].ravel()), (255, 0, 0), thickness) # Z (blue)

        # Ignore errors due to out of field axis: their coordinate are larger than int32 limitations.
        except cv2.error:
            pass

    def draw_places(self, image: numpy.array, K, D, color: tuple = None, border_size: int = 0):
        """Draw group places."""

        l = self.marker_size / 2

        for identifier, place in self.__places.items():

            try:

                placePoints, _ = cv2.projectPoints(place.corners, self.__rotation, self.__translation, numpy.array(K), numpy.array(D))
                placePoints = placePoints.astype(int)
                
                cv2.line(image, tuple(placePoints[0].ravel()), tuple(placePoints[1].ravel()), color, border_size)
                cv2.line(image, tuple(placePoints[1].ravel()), tuple(placePoints[2].ravel()), color, border_size)
                cv2.line(image, tuple(placePoints[2].ravel()), tuple(placePoints[3].ravel()), color, border_size)
                cv2.line(image, tuple(placePoints[3].ravel()), tuple(placePoints[0].ravel()), color, border_size)

            # Ignore errors due to out of field places: their coordinate are larger than int32 limitations.
            except cv2.error:
                pass

    def draw(self, image: numpy.array, K, D, draw_axes: dict = None, draw_places: dict = None):
        """Draw group axes and places.
    
        Parameters:
            draw_axes: draw_axes parameters (if None, no axes drawn)
            draw_places: draw_places parameters  (if None, no places drawn)
        """

        # Draw axes if required
        if draw_axes is not None:

            self.draw_axes(image, K, D, **draw_axes)

        # Draw places if required
        if draw_places is not None:

            self.draw_places(image, K, D, **draw_places)

    def to_obj(self, obj_filepath):
        """Save group to .obj file."""

        with open(obj_filepath, 'w', encoding='utf-8') as file:

            file.write('# ArGaze OBJ File\n')
            file.write('# http://achil.recherche.enac.fr/features/eye/argaze/\n')

            v_count = 0

            for p, (identifier, place) in enumerate(self.__places.items()):

                file.write(f'o {self.__dictionary.name}#{identifier}_Marker\n')

                vertices = ''

                # Write vertices in reverse order
                for v in [3, 2, 1, 0]:

                    file.write(f'v {" ".join(map(str, place.corners[v]))}\n')
                    v_count += 1

                    vertices += f' {v_count}'

                #file.write('s off\n')
                file.write(f'f{vertices}\n')