#!/usr/bin/env python3

import actionlib
import rospy

import yolov5.msg
from sensor_msgs.msg import Image

from classes.yolo_class import YoloDetector

import cv2 # for cv_bridge
from cv_bridge import CvBridge
import numpy as np
import time

class ActionServer():

    def __init__(self):

        t0 = time.time()

        ### various stuff
        self.DO_DEBUGING = True

        ### yolo detector initialization
        device_num = '' # 'cpu' / '' for GPU priority
        self.yolo = YoloDetector(device_num)
        self.yolo.DO_DEBUGING = self.DO_DEBUGING

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

        ### action server initialization
        action_name = 'detectGoal'
        action_definition = yolov5.msg.detectGoalAction
        self.feedback = yolov5.msg.detectGoalFeedback()
        self.result = yolov5.msg.detectGoalResult()
        self.server = actionlib.SimpleActionServer(
            action_name, action_definition, self.execute, False)
        self.server.start()
        self.server.register_preempt_callback(self.setPreemptState)

        ### INFO
        t1 = time.time()
        dt1 = np.round(1000*(t1-t0))
        rospy.loginfo(f'Yolov5 server initialization done in {dt1} ms')


    '''
    Execute callback - action function - goal is action ros msg goal
    '''
    def execute(self, goal):
        
        r = rospy.Rate(5)

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

            ### obtain image 
            imgmsg = rospy.wait_for_message("video_frames", Image)
            img = self.bridge.imgmsg_to_cv2(imgmsg)

            ### perform detection
            self.yolo.perform_detection(img)

            ### determine what was detected and sort detections to flowers and flower pots
            if self.yolo.flag_anything:
                self.yolo.process_detections()

            ### logging, image publishing
            if self.DO_DEBUGING:

                # log results if anything detected
                if self.yolo.flag_any_flower:
   
                    # loginfo
                    conf_str = ''.join(f'{e:.2f}; ' for e in self.yolo.flowers_conf)
                    n_flwrs = len(self.yolo.flowers_conf)
                    rospy.loginfo(f"Detected {n_flwrs} flowers with confidence {conf_str}.")

                    # anotate original image
                    t2 = time.time()
                    self.yolo.anotate_raw()
                    if self.yolo.flag_flower and self.yolo.flag_pot:
                        self.yolo.form_pairs() # TODO pro urcovani polohy; tady jen testuju
                        # anotate pairs also
                        self.yolo.anotate_pairs()

                    t3 = time.time()
                    rospy.loginfo(f'Anotation took {np.round(1000*(t3-t2))} ms')

                # publish image
                self.image_pub.publish(self.bridge.cv2_to_imgmsg(self.yolo.img0, "bgr8"))

            ### if a flower detected with high enough confidence
            # TODO predelat na uspesnou detekci pri detekci paru kytka-kvetinac
            if self.yolo.flag_flower:
                
                ### write results (pose unused)
                self.result.exitcode = 0
                self.result.pose.position.x = 0
                self.result.pose.position.y = 0
                self.result.pose.position.z = 0
                self.result.pose.orientation.x = 0
                self.result.pose.orientation.y = 0
                self.result.pose.orientation.z = 0
                self.result.pose.orientation.w = 0

                self.server.set_succeeded(self.result)     

		    # TODO if something nasty happens
            if False:
                bExecute_detection = False
                ### abort action -> send result to client - problem!
                self.result.exitcode = 1
                self.server.set_succeeded(self.result)
		    
            ### rate of while loop
            r.sleep()
    

    '''
    setPreempt - action go to preempt state, 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) # proc ne set_preempted?


if __name__ == '__main__':
    
    node_name = 'yolov5_server'
    rospy.init_node(node_name)
    server = ActionServer()
    rospy.spin()
