#!/usr/bin/env python3

'''
This file is original version of action_server.py.
It contains all debuging code and is generally messier.
'''

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
from cv_bridge import CvBridge
import cv2
import pathlib
import numpy as np
import tf
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 = 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.bridge = CvBridge()

        ### action server initialization

        # set action name
        action_name = 'detectMarker'
        
        # set action definiton
        action_definition = 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_definition, self.execute, False)
        self.server.start()
        
        # register preempt callback
        self.server.register_preempt_callback(self.setPreemptState)

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

        # start subscribing images from webcam (TODO alternativa: self.imgmsg = rospy.wait_for_message("video_frames",Image))
        #rospy.loginfo("Spousti se subscriber webkamery")
        #self.imgmsg = rospy.wait_for_message("video_frames",Image)
        #self.image_sub = rospy.Subscriber("video_frames", Image, self.webcam_callback)
        
        # for debugging purpose
        self.image_pub = rospy.Publisher("detected_aruco_img",Image, queue_size=1)

        i = 0
        
        ### MAIN LOOP

        while bExecute_detection:

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

            # perform aruco detection and pose estimation + postprocessing
            # markers_img, ids, rvecs, tvecs = self.detect_aruco(self.cv_image, goal)
            img_gray = cv2.cvtColor(self.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 ids is not None:
                if [goal.id] in ids:
                    idx = np.where(ids == goal.id)[0][0]
                    self.id_detected = ids[idx][0]
                
                    rospy.loginfo("Detected ID: %i on index %i.", self.id_detected, idx) # debug

                    ### 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 rotation matrix to Rodrigues params
                    # - not needed, only for visualization, make quaternion straight from this matrix
                    rvec_i = cv2.Rodrigues(rmat_marker)
                    rvec_i = rvec_i[0].reshape(1,-1) # reshape to row vector
                    rvecs[idx] = rvec_i

                    ### convert orinetation to quaternion
                    # using Rodrigues
                    rmat = np.eye(4, dtype=float)
                    rmat[0:3, 0:3], _ = cv2.Rodrigues(rvecs[idx]) # - not needed
                    # rmat[0:3, 0:3] = rmat_marker # better if visualisation is not needed
                    quat = tf.transformations.quaternion_from_matrix(rmat)
                    ## or using rotation vector
                    # angle = cv2.norm(rvecs[idx])
                    # quat = tf.transformations.quaternion_about_axis(angle, rvecs[idx])

                    # toto upravit

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

                    ### draw wanted marker --- debug
                    self.cv_image = cv2.aruco.drawDetectedMarkers(self.cv_image, [corners[idx]], ids[idx])

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

                ### draw wanted marker if detected and all axis crosses --- debug
                for rvec, tvec in zip(rvecs, 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.drawAxis(self.cv_image, self.cam_mtx, \
                #     self.dist_coefs, rvecs[idx], tvecs[idx], self.axis_length)
                # self.cv_image = cv2.aruco.drawDetectedMarkers(self.cv_image, corners, ids)

            ### feedback
            self.feedback.percent_complete = 50.0


            # log detected ids - debug
            if ids is None:
                self.id_detected = None
                rospy.loginfo("Detected IDs: None") # debug
            else:
                #self.id_detected = int(ids[0])
                ids_str = ''.join(str(e) for e in ids)
                rospy.loginfo("Detected IDs: %s.", ids_str) # debug

            ### publishing

            # publish image with detections and poses --- debug
            self.image_pub.publish(self.bridge.cv2_to_imgmsg(self.cv_image, "bgr8"))

            # publish feedback to client
            # self.server.publish_feedback(self.feedback)


            ### debug
            i = i + 1

            # TODO if goal_id is in detected ids
            if i == 15:
                pass

            # TODO if something nasty happens
            if i == 10000:
                bExecute_detection = False
                ### abort action -> send result to client - problem!
                self.server.set_aborted(self.result)
            
            # rate of while loop
            r.sleep()  
    
    '''
    setPreempt - action go to preempt state, usually called when the goal is canceled -> abort action!
    '''
    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)

    '''
    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 = 'my_action_server'
    
    #init node
    rospy.init_node(node_name)
    
    #init server
    server = ActionServer()
    
    rospy.spin()