'''
Class for detecting ArUco markers and estimating their pose
'''

import rospy

from sensor_msgs.msg import Image

import cv2
from cv_bridge import CvBridge
import tf
import numpy as np
from pathlib import Path
import time

FILE = Path(__file__).resolve()
ROOT = FILE.parents[2]

class ArucoDetector:

    def __init__(self, dict, marker_size, calib_path):

        self.DO_DEBUGING = True

        ### aruco initialization

        ARUCO_DICTIONARY = {
            'DICT_ARUCO_ORIGINAL': cv2.aruco.DICT_ARUCO_ORIGINAL,
            'DICT_5X5_1000': cv2.aruco.DICT_5X5_1000,
        }

        # load camera calibration files
        calib_files = np.load(ROOT / calib_path)
        self.cam_mtx = calib_files['cam_mtx']
        self.dist_coefs = calib_files['dist_coefs']

        # init aruco utilities                   
        self.aruco_dict = cv2.aruco.getPredefinedDictionary(ARUCO_DICTIONARY[dict])
        self.marker_length = marker_size # [m]
        self.axis_length = self.marker_length/2 # [m] for display of axis cross
        self.arucoParams = cv2.aruco.DetectorParameters_create()
        self.arucoParams.minMarkerPerimeterRate = 0.03 # 0.03
        self.arucoParams.maxMarkerPerimeterRate = 4.0 # 4.0

        # init atributes
        self.ids = None
        self.tvecs = None
        self.rvecs = None
        self.id_detected = None
        self.tvec = None
        self.rvec = None

        ### ROS initalization
        self.bridge = CvBridge()
        self.image_sub = rospy.Subscriber("video_frames", Image)
        if self.DO_DEBUGING:
            self.image_pub = rospy.Publisher("detected_aruco_img", Image, queue_size=1)

    '''
    obtain image, detect markers and estimate their poses
    '''
    def perform_detection(self):
        

        # get image
        imgmsg = rospy.wait_for_message("video_frames",Image)
        self.cv_image = self.bridge.imgmsg_to_cv2(imgmsg)
        img_gray = cv2.cvtColor(self.cv_image, cv2.COLOR_BGR2GRAY)

        t1 = time.time()

        # detect markers
        self.corners, self.ids, rejected = cv2.aruco.detectMarkers(
            img_gray, self.aruco_dict, parameters=self.arucoParams)

        # estimate pose of detected markers
        self.rvecs, self.tvecs, _ = cv2.aruco.estimatePoseSingleMarkers(
            self.corners, self.marker_length, self.cam_mtx, self.dist_coefs)

        t2 = time.time()
        
        if self.DO_DEBUGING:
            rospy.loginfo(f'Detection time: {np.round(1000*(t2-t1))} ms')

    '''
    Process detections of all markers to return filtered pose of wanted marker
    '''
    def process_detections(self):

        ### info about wanted maker
        idx = np.nonzero(self.ids == self.goal_id)[0][0]
        self.id_detected = self.ids[idx][0]
        self.tvec = self.tvecs[idx][0]
        self.rvec = self.rvecs[idx][0]

        ### rotate marker default CS to match camera CS
        
        # convert Rodrigues params to rotation matrix
        rmat_marker = cv2.Rodrigues(self.rvec)[0]
        # rotate marker about x axis by 180°
        rmat_x180 = np.eye(3, dtype=float) * -1
        rmat_x180[0,0] = 1
        rmat_marker = np.matmul(rmat_marker, rmat_x180)

        ### convert orinetation to quaternion
        rmat = np.eye(4, dtype=float)
        rmat[0:3, 0:3] = rmat_marker # rotated orientation
        self.quat = tf.transformations.quaternion_from_matrix(rmat)

    '''
    Write to result
    '''
    def write_results(self):
        self.result.id_detected = self.id_detected
        self.result.exitcode = 0
        self.result.pose.position.x = self.tvec[0]
        self.result.pose.position.y = self.tvec[1]
        self.result.pose.position.z = self.tvec[2]
        self.result.pose.orientation.x = self.quat[0]
        self.result.pose.orientation.y = self.quat[1]
        self.result.pose.orientation.z = self.quat[2]
        self.result.pose.orientation.w = self.quat[3]

    '''
    Write to feedback - aruco continuous
    '''
    def write_feedback(self):
        self.feedback.id_detected = self.id_detected
        self.result.exitcode = 0
        self.feedback.pose.position.x = self.tvec[0]
        self.feedback.pose.position.y = self.tvec[1]
        self.feedback.pose.position.z = self.tvec[2]
        self.feedback.pose.orientation.x = self.quat[0]
        self.feedback.pose.orientation.y = self.quat[1]
        self.feedback.pose.orientation.z = self.quat[2]
        self.feedback.pose.orientation.w = self.quat[3]

    '''
    Draw ids and axis crosses of detected markers - for DEBUGGING
    '''
    def draw_detections(self):
        for rvec, tvec in zip(self.rvecs, self.tvecs):
            self.cv_image = cv2.aruco.drawAxis(self.cv_image, self.cam_mtx,
                self.dist_coefs, rvec, tvec, self.axis_length)
        self.cv_image = cv2.aruco.drawDetectedMarkers(self.cv_image, self.corners, self.ids)

        