#!/usr/bin/env python

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

from argaze.ArUcoMarkers import ArUcoMarkersDictionary, ArUcoMarker, ArUcoCamera

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

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."""

ArUcoSceneType = TypeVar('ArUcoScene', bound="ArUcoScene")
# 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))

@dataclass(frozen=True)
class Place():
    """Define a place as a pose and a marker."""

    translation: numpy.array
    """Position in scene referential."""

    rotation: numpy.array
    """Rotation in scene referential."""

    marker: dict
    """ArUco marker linked to the place."""

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

    marker_size: float = field(default=0.)
    """Expected size of all markers in the scene."""

    dictionary: ArUcoMarkersDictionary.ArUcoMarkersDictionary = field(default_factory=ArUcoMarkersDictionary.ArUcoMarkersDictionary)
    """Expected dictionary of all markers in the scene."""

    places: dict = field(default_factory=dict)
    """Expected markers place"""

    def __post_init__(self):
        """Initialize cached data to speed up further processings."""

        # Init pose data
        self._translation = numpy.zeros(3)
        self._rotation = numpy.zeros(3)

        # Normalize places data
        new_places = {}

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

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

                identifier = int(identifier)

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

                # Check rotation value shape
                rvalue = numpy.array(place.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}')

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

                new_places[identifier] = Place(tvec, rmat, new_marker)

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

                new_places[identifier] = place

        self.places = new_places

        # Process axis-angle between place combination to speed up further calculations
        self.__angle_cache = {}
        for (A_identifier, A_place), (B_identifier, B_place) in itertools.combinations(self.places.items(), 2):

            A = self.places[A_identifier].rotation
            B = self.places[B_identifier].rotation

            if numpy.array_equal(A, B):

                angle = 0.

            else:

                # Rotation matrix from A place to B place
                AB = B.dot(A.T)

                # Calculate axis-angle representation of AB rotation matrix
                angle = numpy.rad2deg(numpy.arccos((numpy.trace(AB) - 1) / 2))

            try:
                self.__angle_cache[A_identifier][B_identifier] = angle
            except:
                self.__angle_cache[A_identifier] = {B_identifier: angle}

            try:
                self.__angle_cache[B_identifier][A_identifier] = angle
            except:
                self.__angle_cache[B_identifier] = {A_identifier: angle}

        # Process distance between each place combination to speed up further calculations
        self.__distance_cache = {}
        for (A_identifier, A_place), (B_identifier, B_place) in itertools.combinations(self.places.items(), 2):

            A = self.places[A_identifier].translation
            B = self.places[B_identifier].translation

            # Calculate axis-angle representation of AB rotation matrix
            distance = numpy.linalg.norm(B - A)

            try:
                self.__distance_cache[A_identifier][B_identifier] = distance
            except:
                self.__distance_cache[A_identifier] = {B_identifier: distance}

            try:
                self.__distance_cache[B_identifier][A_identifier] = distance
            except:
                self.__distance_cache[B_identifier] = {A_identifier: distance}

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

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

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

        .. note::
           Marker normal vectors (vn) expected.

        """

        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'),
            'normal': re.compile(r'vn ([+-]?[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]+) ([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 = []
            normals = {}
            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 scene dictionary with first dictionary name
                        if new_dictionary == None:

                            new_dictionary = ArUcoMarkersDictionary.ArUcoMarkersDictionary(dictionary)

                        # Check all others marker dictionary are equal to new scene 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 normal to calculate rotation matrix
                    elif key == 'normal':

                        normals[identifier] = 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(3)), int(match.group(5)), int(match.group(7))]

                    # 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 from counter clockwise ordered face vertices
                    corners = numpy.array([ vertices[i-1] for i in face ])

                    # Edit translation (Tp) allowing to move world axis (W) at place axis (P)
                    Tp = corners.mean(axis=0)

                    # Edit place axis from corners positions
                    place_x_axis = corners[1:3].mean(axis=0) - Tp
                    place_x_axis_norm = numpy.linalg.norm(place_x_axis)
                    place_x_axis = place_x_axis / place_x_axis_norm

                    place_y_axis = corners[2:4].mean(axis=0) - Tp
                    place_y_axis_norm = numpy.linalg.norm(place_y_axis)
                    place_y_axis = place_y_axis / place_y_axis_norm

                    place_z_axis = normals[identifier]

                    # Edit rotation (Rp) allowing to transform world axis (W) into place axis (P)
                    W = numpy.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
                    P = numpy.array([place_x_axis, place_y_axis, place_z_axis])
                    Rp = W.dot(P.T)

                    # 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*2

                        # 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(Tp, Rp, new_marker)

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

        return ArUcoScene(new_marker_size, new_dictionary, new_places)

    @classmethod
    def from_json(self, json_filepath: str) -> ArUcoSceneType:
        """Load ArUco scene 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 ArUcoScene(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.translation}'
            output += f'\n{place.rotation}'

        output += '\n\n\tAngle cache:'
        for A_identifier, A_angle_cache in self.__angle_cache.items():
            for B_identifier, angle in A_angle_cache.items():
                output += f'\n\t\t- {A_identifier}/{B_identifier}: {angle:3f}'

        output += '\n\n\tDistance cache:'
        for A_identifier, A_distance_cache in self.__distance_cache.items():
            for B_identifier, distance in A_distance_cache.items():
                output += f'\n\t\t- {A_identifier}/{B_identifier}: {distance:3f}'

        return output

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

        return list(self.places.keys())

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

        * **Returns:**
            - dict of markers belonging to this scene
            - dict of remaining markers not belonging to this scene
        """

        scene_markers = {}
        remaining_markers = {}

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

            if marker_id in self.places.keys():

                scene_markers[marker_id] = marker

            else:
                
                remaining_markers[marker_id] = marker

        return scene_markers, remaining_markers

    def check_markers_consistency(self, scene_markers: dict, angle_tolerance: float, distance_tolerance: float) -> Tuple[dict, dict]:
        """Evaluate if given markers configuration match related places configuration.

        * **Returns:**
            - dict of consistent markers
            - dict of unconsistent markers
            - dict of identified distance or angle unconsistencies and out-of-bounds values
        """

        consistent_markers = {}
        unconsistencies = {}

        for (A_identifier, A_marker), (B_identifier, B_marker) in itertools.combinations(scene_markers.items(), 2):

            try:

                # Rotation matrix from A marker to B marker
                AB = B_marker.rotation.dot(A_marker.rotation.T)

                # Calculate axis-angles representation of AB rotation matrix
                angle = numpy.rad2deg(numpy.arccos((numpy.trace(AB) - 1) / 2))
                expected_angle = self.__angle_cache[A_identifier][B_identifier]

                # Calculate distance between A marker center and B marker center
                distance = numpy.linalg.norm(A_marker.translation - B_marker.translation)
                expected_distance = self.__distance_cache[A_identifier][B_identifier]

                # Check angle and distance according given tolerance then normalise marker pose
                consistent_angle = math.isclose(angle, expected_angle, abs_tol=angle_tolerance)
                consistent_distance = math.isclose(distance, expected_distance, abs_tol=distance_tolerance)

                if consistent_angle and consistent_distance:

                    if A_identifier not in consistent_markers.keys():

                        # Remember this marker is already validated
                        consistent_markers[A_identifier] = A_marker

                    if B_identifier not in consistent_markers.keys():

                        # Remember this marker is already validated
                        consistent_markers[B_identifier] = B_marker

                else:

                    if not consistent_angle:
                        unconsistencies[f'{A_identifier}/{B_identifier} angle'] = angle
                    
                    if not consistent_distance:
                        unconsistencies[f'{A_identifier}/{B_identifier} distance'] = distance

            except KeyError:

                raise ValueError(f'Marker {A_identifier} or {B_identifier} don\'t belong to the scene.')

        # Gather unconsistent markers
        unconsistent_markers = {}

        for name, marker in scene_markers.items():

            if name not in consistent_markers:

                unconsistent_markers[name] = marker

        return consistent_markers, unconsistent_markers, unconsistencies

    def estimate_pose_from_single_marker(self, marker: ArUcoMarker.ArUcoMarker) -> Tuple[numpy.array, numpy.array]:
        """Calculate rotation and translation that move a marker to its place."""

        # Get the place related to the given marker
        try:

            place = self.places[marker.identifier]

            # Rotation matrix that transform marker to related place
            self._rotation = marker.rotation.dot(place.rotation.T)

            # Translation vector that transform marker to related place
            self._translation = marker.translation - place.translation.dot(place.rotation).dot(marker.rotation.T)

            return self._translation, self._rotation

        except KeyError:

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

    def estimate_pose_from_markers(self, markers: dict) -> Tuple[numpy.array, numpy.array]:
        """Calculate average rotation and translation that move markers to their related places."""

        rotations = []
        translations = []

        for identifier, marker in markers.items():

            try:

                place = self.places[identifier]

                # Rotation matrix that transform marker to related place
                R = marker.rotation.dot(place.rotation.T)

                # Translation vector that transform marker to related place
                T = marker.translation - place.translation.dot(place.rotation).dot(marker.rotation.T)

                rotations.append(R)
                translations.append(T)

            except KeyError:

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

        # Consider ArUcoScene rotation as the mean of all marker rotations
        # !!! WARNING !!! This is a bad hack : processing rotations average is a very complex problem that needs to well define the distance calculation method before.
        self._rotation = numpy.mean(numpy.array(rotations), axis=0)

        # Consider ArUcoScene translation as the mean of all marker translations
        self._translation = numpy.mean(numpy.array(translations), axis=0)

        return self._translation, self._rotation

    def estimate_pose_from_axis_markers(self, origin_marker: ArUcoMarker.ArUcoMarker, horizontal_axis_marker: ArUcoMarker.ArUcoMarker, vertical_axis_marker: ArUcoMarker.ArUcoMarker) -> Tuple[numpy.array, numpy.array]:
        """Calculate rotation and translation from 3 markers defining an orthogonal axis."""

        O_marker = origin_marker
        A_marker = horizontal_axis_marker
        B_marker = vertical_axis_marker

        O_place = self.places[O_marker.identifier]
        A_place = self.places[A_marker.identifier]
        B_place = self.places[B_marker.identifier]

        # Place axis
        OA = A_place.translation - O_place.translation
        OA = OA / numpy.linalg.norm(OA)

        OB = B_place.translation - O_place.translation
        OB = OB / numpy.linalg.norm(OB)

        # Detect and correct bad place axis orientation
        X_sign = numpy.sign(OA)[0]
        Y_sign = numpy.sign(OB)[1]

        P = numpy.array([OA*X_sign, OB*Y_sign, numpy.cross(OA*X_sign, OB*Y_sign)])

        # Marker axis
        OA = A_marker.translation - O_marker.translation
        OA = OA / numpy.linalg.norm(OA)

        OB = B_marker.translation - O_marker.translation
        OB = OB / numpy.linalg.norm(OB)

        # Detect and correct bad place axis orientation
        X_sign = numpy.sign(OA)[0]
        Y_sign = -numpy.sign(OB)[1]

        M = numpy.array([OA*X_sign, OB*Y_sign, numpy.cross(OA*X_sign, OB*Y_sign)])

        # Then estimate ArUcoScene rotation
        self._rotation = P.dot(M.T)

        # Consider ArUcoScene translation as the translation of the marker at axis origin
        self._translation = O_marker.translation - O_place.translation.dot(O_place.rotation).dot(M.T)

        return self._translation, self._rotation

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

        return self._translation

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

        self._translation = tvec

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

        return self._translation

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

        self._rotation = rmat

    def draw_axis(self, frame, K, D, consistency=2):
        """Draw scene axis according a consistency score."""

        l = self.marker_size / 2
        ll = self.marker_size

        # Select color according consistency score
        n = 95 * consistency if consistency < 2 else 0
        f = 159 * consistency if consistency < 2 else 255

        # Draw axis
        axisPoints = numpy.float32([[ll, 0, 0], [0, ll, 0], [0, 0, ll], [0, 0, 0]]).reshape(-1, 3)
        axisPoints, _ = cv.projectPoints(axisPoints, self._rotation, self._translation, numpy.array(K), numpy.array(D))
        axisPoints = axisPoints.astype(int)

        cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[0].ravel()), (n,n,f), 6) # X (red)
        cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[1].ravel()), (n,f,n), 6) # Y (green)
        cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[2].ravel()), (f,n,n), 6) # Z (blue)

    def draw_places(self, frame, K, D, consistency=2):
        """Draw scene places and their axis according a consistency score."""

        l = self.marker_size / 2
        ll = self.marker_size

        # Select color according consistency score
        n = 95 * consistency if consistency < 2 else 0
        f = 159 * consistency if consistency < 2 else 255

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

            try:

                T = self.places[identifier].translation
                R = self.places[identifier].rotation

                # Draw place axis
                axisPoints = (T + numpy.float32([R.dot([l/2, 0, 0]), R.dot([0, l/2, 0]), R.dot([0, 0, l/2]), R.dot([0, 0, 0])])).reshape(-1, 3)
                axisPoints, _ = cv.projectPoints(axisPoints, self._rotation, self._translation, numpy.array(K), numpy.array(D))
                axisPoints = axisPoints.astype(int)

                cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[0].ravel()), (n,n,f), 6) # X (red)
                cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[1].ravel()), (n,f,n), 6) # Y (green)
                cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[2].ravel()), (f,n,n), 6) # Z (blue)
                
                # Draw place
                placePoints = (T + numpy.float32([R.dot([-l, -l, 0]), R.dot([l, -l, 0]), R.dot([l, l, 0]), R.dot([-l, l, 0])])).reshape(-1, 3)
                placePoints, _ = cv.projectPoints(placePoints, self._rotation, self._translation, numpy.array(K), numpy.array(D))
                placePoints = placePoints.astype(int)
                
                cv.line(frame, tuple(placePoints[0].ravel()), tuple(placePoints[1].ravel()), (f,f,f), 3)
                cv.line(frame, tuple(placePoints[1].ravel()), tuple(placePoints[2].ravel()), (f,f,f), 3)
                cv.line(frame, tuple(placePoints[2].ravel()), tuple(placePoints[3].ravel()), (f,f,f), 3)
                cv.line(frame, tuple(placePoints[3].ravel()), tuple(placePoints[0].ravel()), (f,f,f), 3)

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