#!/usr/bin/env python

from typing import TypeVar, Tuple
from dataclasses import dataclass, field
import json
import os

from argaze.ArUcoMarkers import *
from argaze.AreaOfInterest import *

import numpy

ArSceneType = TypeVar('ArScene', bound="ArScene")
# Type definition for type annotation convenience

AOI2DSceneType = TypeVar('AOI2DScene', bound="AOI2DScene")
# Type definition for type annotation convenience

class PoseEstimationFailed(Exception):
	"""Exception raised by ArScene project method when the pose can't be estimated due to unconsistencies."""

	def __init__(self, message, unconsistencies=None):  

		super().__init__(message)
        
		self.unconsistencies = unconsistencies

class SceneProjectionFailed(Exception):
	"""Exception raised by ArScene project method when the scene can't be projected."""

	def __init__(self, message):  

		super().__init__(message)

@dataclass
class ArScene():
	"""Define an Augmented Reality environnement thanks to ArUco markers and project it onto incoming frames."""

	name: str 
	"""Project name."""

	aruco_dictionary: ArUcoMarkersDictionary.ArUcoMarkersDictionary = field(init=False, default_factory=ArUcoMarkersDictionary.ArUcoMarkersDictionary)
	"""ArUco markers dictionary."""

	aruco_marker_size: float = field(init=False)
	"""Size of ArUco markers in centimeter."""

	aruco_camera: ArUcoCamera.ArUcoCamera = field(init=False, default_factory=ArUcoCamera.ArUcoCamera)
	"""ArUco camera ..."""

	aruco_tracker: ArUcoTracker.ArUcoTracker = field(init=False, default_factory=ArUcoTracker.ArUcoTracker)
	"""ArUco tracker ..."""

	aruco_scene: ArUcoScene.ArUcoScene = field(init=False, default_factory=ArUcoScene.ArUcoScene)
	"""ArUco scene ..."""

	aoi_scene: AOI3DScene.AOI3DScene = field(init=False, default_factory=AOI3DScene.AOI3DScene)
	"""AOI 3D scene ..."""

	angle_tolerance: float
	"""Angle error tolerance allowed to validate marker pose in degree."""

	distance_tolerance: float
	"""Distance error tolerance allowed to validate marker pose in centimeter."""

	aruco_axis: dict
	"""Dictionary of orthogonal axis where each axis is defined by list of 3 markers identifier (first is origin)."""

	aruco_aoi: dict
	"""Dictionary of AOI defined by list of markers identifier and markers corners index tuples."""

	def __init__(self, **kwargs):

		self.aruco_dictionary = ArUcoMarkersDictionary.ArUcoMarkersDictionary(kwargs.pop('aruco_dictionary'))

		self.aruco_marker_size = kwargs.pop('aruco_marker_size')

		self.aruco_camera = ArUcoCamera.ArUcoCamera(**kwargs.pop('aruco_camera'))

		self.aruco_tracker = ArUcoTracker.ArUcoTracker(self.aruco_dictionary, self.aruco_marker_size, self.aruco_camera, **kwargs.pop('aruco_tracker'))

		# Check aruco_scene value type
		aruco_scene_value = kwargs.pop('aruco_scene')

		# str: relative path to .obj file
		if type(aruco_scene_value) == str:

			aruco_scene_value = os.path.join(self.__current_directory, aruco_scene_value)

		self.aruco_scene = ArUcoScene.ArUcoScene(self.aruco_dictionary, self.aruco_marker_size, aruco_scene_value)

		# Check aoi_scene value type
		aoi_scene_value = kwargs.pop('aoi_scene')

		# str: relative path to .obj file
		if type(aoi_scene_value) == str:

			obj_filepath = os.path.join(self.__current_directory, aoi_scene_value)
			self.aoi_scene = AOI3DScene.AOI3DScene.from_obj(obj_filepath)

		# dict: all AOI
		else:
			self.aoi_scene = AOI3DScene.AOI3DScene(aoi_scene_value)

		# Init aruco axis
		self.aruco_axis = {}

		# Init aruco aoi
		self.aruco_aoi = {}

		# Update all attributes from arguments
		self.__dict__.update(kwargs)

		# Convert aruco axis markers identifier into ARUCO_DICT_NAME#ID string
		aruco_axis_string = {}
		for axis_name, markers_id in self.aruco_axis.items():

			# Estimate pose from axis markers
			aruco_axis_names = []
			for marker_id in markers_id:
				aruco_axis_names.append(f'{self.aruco_dictionary.name}#{marker_id}')

			aruco_axis_string[axis_name] = aruco_axis_names

		self.aruco_axis = aruco_axis_string

	@classmethod
	def from_json(self, json_filepath: str) -> ArSceneType:
		"""Load ArGaze project from .json file."""

		with open(json_filepath) as configuration_file:

			# Store current directory to allow relative path loading
			self.__current_directory = os.path.dirname(os.path.abspath(json_filepath))
        
			return ArScene(**json.load(configuration_file))

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

		output = ''
		output += f'\nArUcoCamera: {self.aruco_camera}'
		output += f'\n\nArUcoTracker tracking data: {self.aruco_tracker.tracking_data}'
		output += f'\n\nArUcoScene: {self.aruco_scene}'
		output += f'\n\nAOIScene: {self.aoi_scene}'
		output += '\n'

		return output

	def estimate_pose(self, frame) -> Tuple[numpy.array, numpy.array, dict]:
		"""Estimate scene pose from ArUco markers into frame.

		* **Returns:** 
			- scene translation vector
			- scene rotation matrix
			- dict of markers used to estimate the pose
		"""

		self.aruco_tracker.track(frame)

		# Pose estimation fails when no marker is detected
		if len(self.aruco_tracker.tracked_markers) == 0:
			
			raise PoseEstimationFailed('No marker detected')

		scene_markers, _ = self.aruco_scene.filter_markers(self.aruco_tracker.tracked_markers)

		# Pose estimation fails when no marker belongs to the scene
		if len(scene_markers) == 0:

			raise PoseEstimationFailed('No marker belongs to the scene')

		# Estimate scene pose from unique marker transformations
		elif len(scene_markers) == 1:

			tvec, rmat = self.aruco_scene.estimate_pose_from_any_markers(scene_markers)

			return tvec, rmat, scene_markers

		# Try to estimate scene pose from 3 markers defining an orthogonal axis
		elif len(scene_markers) >= 3 and len(self.aruco_axis) > 0:

			for axis_name, markers_names in self.aruco_axis.items():

				try:

					axis_markers = []
					for name in markers_names:
						axis_markers.append((name, scene_markers[name]))

					tvec, rmat = self.aruco_scene.estimate_pose_from_axis_markers(axis_markers)

					return tvec, rmat, axis_markers

				except:
					pass

			raise PoseEstimationFailed('No marker axis')

		# Otherwise, check markers consistency
		consistent_markers, unconsistent_markers, unconsistencies = self.aruco_scene.check_markers_consistency(scene_markers, self.angle_tolerance, self.distance_tolerance)

		# Pose estimation fails when no marker passes consistency checking
		if len(consistent_markers) == 0:

			raise PoseEstimationFailed('Unconsistent marker poses', unconsistencies)

		# Otherwise, estimate scene pose from all markers transformations
		tvec, rmat = self.aruco_scene.estimate_pose_from_any_markers(consistent_markers)

		return tvec, rmat, consistent_markers

	def project(self, tvec, rmat, visual_hfov=0) -> AOI2DSceneType:
		"""Project AOI scene into frame according estimated pose."""

		# Clip AOI out of the visual horizontal field of view (optional)
		if visual_hfov > 0:

			# Transform scene into camera referential
			aoi_scene_camera_ref = self.aoi_scene.transform(tvec, rmat)

			# Get aoi inside vision cone field 
			cone_vision_height_cm = 200 # cm
			cone_vision_radius_cm = numpy.tan(numpy.deg2rad(visual_hfov / 2)) * cone_vision_height_cm

			_, aoi_outside = aoi_scene_camera_ref.vision_cone(cone_vision_radius_cm, cone_vision_height_cm)

			# Keep only aoi inside vision cone field
			aoi_scene_copy = self.aoi_scene.copy(exclude=aoi_outside.keys())

		else:

			aoi_scene_copy = self.aoi_scene.copy()

		# DON'T APPLY CAMERA DISTORSION : it projects points which are far from the frame into it
		# This hack isn't realistic but as the gaze will mainly focus on centered AOI, where the distorsion is low, it is acceptable.
		aoi_scene_projection = aoi_scene_copy.project(tvec, rmat, self.aruco_camera.K)

		# Warn user when the projected scene is empty
		if len(aoi_scene_projection) == 0:

			raise SceneProjectionFailed('AOI projection is empty')

		return aoi_scene_projection

	def project_aruco_aoi(self, frame) -> AOI2DSceneType:
		"""Edit AOI scene from ArUco markers into frame as defined in aruco_aoi dictionary."""

		self.aruco_tracker.track(frame, estimate_pose=False)

		# AOI projection fails when no marker is detected
		if len(self.aruco_tracker.tracked_markers) == 0:
			
			raise SceneProjectionFailed('No marker detected')

		scene_markers, _ = self.aruco_scene.filter_markers(self.aruco_tracker.tracked_markers)

		# AOI projection fails when no marker belongs to the scene
		if len(scene_markers) == 0:

			raise SceneProjectionFailed('No marker belongs to the scene')

		aoi_scene = {}

		for name, marker_corners in self.aruco_aoi.items():

			aoi_points = []
			for marker_id, corner_id in marker_corners:

				aoi_points.append(self.aruco_tracker.tracked_markers[marker_id].corners[0][corner_id])

			aoi_scene[name] = AOIFeatures.AreaOfInterest(aoi_points)

		return AOI2DScene.AOI2DScene(aoi_scene)

	def draw_axis(self, frame):
		"""Draw scene axis into frame."""

		self.aruco_scene.draw_axis(frame, self.aruco_camera.K, self.aruco_camera.D)

	def draw_places(self, frame):
		"""Draw scene places into frame."""

		self.aruco_scene.draw_places(frame, self.aruco_camera.K, self.aruco_camera.D)