#!/usr/bin/env python3

import actionlib
import rospy

# import msgs
import aruco.msg

# imports for aruco detection
from std_msgs.msg import String
from sensor_msgs.msg import Image
import cv2
from cv_bridge import CvBridge
import tf
import numpy as np
from pathlib import Path

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

class ActionServer:
    def __init__(self):

        ### launch parameters
        try:
            DICT = str(rospy.get_param('aruco_dict'))
            MARKER_SIZE = float(rospy.get_param('marker_size'))
            CALIB_PATH = str(rospy.get_param('calib_path'))
        except Exception:
            print("Can't load Aruco parameters from ROS parameter server")
            DICT = 'DICT_ARUCO_ORIGINAL'
            MARKER_SIZE = 0.077
            CALIB_PATH = 'data/calib_params_MS_1280_800.npz'

        ### 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_path = ROOT / CALIB_PATH
        calib_files = np.load(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

        ### ROS initalization
        self.image_sub = rospy.Subscriber("video_frames", Image)

        self.bridge = CvBridge()

        ### action server initialization

        # set action name
        action_name = 'detectMarker'
        
        # set action definiton
        action_def = aruco.msg.detectMarkerAction
        
        # set action feedback and result
        self.feedback = aruco.msg.detectMarkerFeedback()
        self.result = aruco.msg.detectMarkerResult()
        
        # create action server and start it
        self.server = actionlib.SimpleActionServer(action_name, action_def, self.execute, False)
        self.server.start()
        
        # register preempt callback
        self.server.register_preempt_callback(self.setPreemptState)

        # INFO
        rospy.loginfo("Aruco server initialization done.")


    '''
    Execute callback - action function - goal is action ros msg goal
    '''
    def execute(self, goal):
        
        ### INITIALIZATION

        # define framerate of aruco detection (vhodne cca 5-20 Hz)
        r = rospy.Rate(10)
        bExecute_detection = True

        ### MAIN LOOP
        while bExecute_detection and not rospy.is_shutdown() and self.server.is_active():

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

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

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

            # process detections if the wanted id has been found
            if ids is not None:
                if [goal.id] in ids:
                    idx = np.where(ids == goal.id)[0][0]
                    self.id_detected = ids[idx][0]

                    ### rotate marker default CS to match camera CS
                    
                    # convert Rodrigues params to rotation matrix
                    rmat_marker = cv2.Rodrigues(rvecs[idx])[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
                    quat = tf.transformations.quaternion_from_matrix(rmat)

                    ### write results
                    self.result.id_detected = self.id_detected
                    self.result.exitcode = 0
                    self.result.pose.position.x = tvecs[idx][0][0]
                    self.result.pose.position.y = tvecs[idx][0][1]
                    self.result.pose.position.z = tvecs[idx][0][2]
                    self.result.pose.orientation.x = quat[0]
                    self.result.pose.orientation.y = quat[1]
                    self.result.pose.orientation.z = quat[2]
                    self.result.pose.orientation.w = quat[3]

                    ### publish result to client - ID found successfully!
                    bExecute_detection = False
                    self.server.set_succeeded(self.result)

                else:
                    # wanted id not detected yet
                    self.id_detected = None

            # ### feedback
            # self.feedback.percent_complete = 50.0
            # # publish feedback to client
            # self.server.publish_feedback(self.feedback)

            # TODO if something nasty happens
            if False:
                bExecute_detection = False
                ### abort action -> send result to client - problem!
                self.server.set_succeeded(self.result)
            
            # rate of while loop
            r.sleep()  
    
    '''
    setPreempt - action go to preempt state, usually called when the goal is canceled
    '''
    def setPreemptState(self):
        # TODO tidy up the preempted action, then abort the action and send the result
        print('Preempt action')
        # stop subscribing images
        self.image_sub.unregister()

        self.result.exitcode = 10
        self.server.set_succeeded(self.result)

    ''' # NOT USED NOW
    webcam_callback - called when a webcam image is subscribed
    '''
    def webcam_callback(self,data):
        # obtain image
        self.cv_image = self.bridge.imgmsg_to_cv2(data)

if __name__ == '__main__':
    # def node parameters
    node_name = 'aruco_server'
    
    # init node
    rospy.init_node(node_name)
    
    # init server
    server = ActionServer()
    
    rospy.spin()
