#!/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.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():
    """Handle group of ArUco markers as one unique spatial entity and estimate its pose.

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

    marker_size: float = field(default=0.)
    dictionary: ArUcoMarkersDictionary.ArUcoMarkersDictionary = field(default_factory=ArUcoMarkersDictionary.ArUcoMarkersDictionary)
    places: dict = field(default_factory=dict)

    def __post_init__(self):
        """Init group pose and places pose."""

        # Init pose data
        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))

                new_marker = ArUcoMarker.ArUcoMarker(self.dictionary, identifier, self.marker_size)

                # Build marker corners thanks to translation vector and rotation matrix
                place_corners = numpy.array([[-self.marker_size/2, self.marker_size/2, 0], [self.marker_size/2, self.marker_size/2, 0], [self.marker_size/2, -self.marker_size/2, 0], [-self.marker_size/2, -self.marker_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

    @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_marker_size = 0
        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):

                        current_marker_size = place_x_axis_norm

                        # Check that all markers size are almost equal
                        if new_marker_size > 0:

                            if not math.isclose(current_marker_size, new_marker_size, rel_tol=1e-3):

                                raise ValueError('Markers size should be almost equal.')

                        new_marker_size = current_marker_size

                    # 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_marker_size, new_dictionary, new_places)

    @classmethod
    def from_json(self, json_filepath: str) -> ArUcoMarkersGroupType:
        """Load ArUco markers group from .json file."""

        new_marker_size = 0
        new_dictionary = None
        new_places = {}

        with open(json_filepath) as configuration_file:

            data = json.load(configuration_file)

            new_marker_size = data.pop('marker_size')
            new_dictionary = data.pop('dictionary')
            new_places = data.pop('places')

        return ArUcoMarkersGroup(new_marker_size, new_dictionary, new_places)

    def __str__(self) -> str:
        """String display"""

        output = f'\n\tDictionary: {self.dictionary}'

        output += f'\n\tMarker size: {self.marker_size} cm'

        output += '\n\n\tPlaces:'
        for identifier, place in self.places.items():
            output += f'\n\t\t- {identifier}:'
            output += f'\n{place.corners}'

        return output

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

        return list(self.places.keys())

    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

    @property
    def translation(self) -> numpy.array:
        """Access to group translation vector."""

        return self._translation

    @translation.setter
    def translation(self, tvec):

        self._translation = tvec

    @property
    def rotation(self) -> numpy.array:
        """Access to group rotation matrix."""

        return self._translation

    @rotation.setter
    def rotation(self, rmat):

        self._rotation = rmat

    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')