#!/usr/bin/env python3

# this script replacates experiments bz loading saved images and lidar data

import actionlib
import rospy

import yolov5.msg
from sensor_msgs.msg import Image, LaserScan

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

import pickle
from datetime import datetime # for experimental data saving

from classes.yolo_class import YoloDetector
from classes.midas_class import MidasDepth
from classes.fusion_class import DepthFusion

class ActionServer():

    def __init__(self):

        t0 = time.time()

        ### various stuff
        self.DO_DEBUGING = True
        self.DATA_LOADING = True # load image and lidar data from file
        self.DATA_SAVING = True # log data about detections and position

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

        ### midas depth estimator initialization
        midas_device_num = 'cpu' # 'cpu' / '' - '' for GPU priority
        self.midas = MidasDepth(midas_device_num)
        self.midas.DO_DEBUGING = self.DO_DEBUGING

        ### lidar & depth map fuser initialization
        self.depth_fusor = DepthFusion()
        self.depth_fusor.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 = 'objectPosition'
        action_definition = yolov5.msg.objectPositionAction
        self.feedback = yolov5.msg.objectPositionFeedback()
        self.result = yolov5.msg.objectPositionResult()
        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 position server initialization done in {dt1} ms')


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

        FILE = Path(__file__).resolve()
        ROOT = FILE.parents[0]
        log_dir = ROOT/'logging_Breach'
        save_dir = ROOT/'logging_2'

        ### MAIN LOOP
        # while not rospy.is_shutdown() and self.server.is_active():
        ######### LOAD MEASURED DATA ######################
        for src_path in log_dir.rglob('*img_original.png'):
            # [0:19] Leela, [0:26] Breach
            name = src_path.name[0:26] # date and time (no decimals)

            
            ### --- LIDAR SCAN ACQUISITION --- ###

            if not self.DATA_LOADING:
                lidar_msg = rospy.wait_for_message("/scan", LaserScan)
            else:
                with open(log_dir / f'{name}_lidar.pkl', 'rb') as f:
                    lidar_msg = pickle.load(f)

            ### --- IMAGE ACQUISITION --- ###

            if not self.DATA_LOADING:
                imgmsg = rospy.wait_for_message("video_frames", Image)
                img = self.bridge.imgmsg_to_cv2(imgmsg)
            else:
                img = cv2.imread(str(log_dir / f'{name}_img_original.png'))
            
            ### --- YOLO DETECTION --- ###

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

            ### try to form pairs
            if self.yolo.flag_flower and self.yolo.flag_pot:
                self.yolo.form_pairs()

            ### logging
            if self.yolo.flag_any_flower:
                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}.")


            ### if flower-pot pairs detected, estimate 3D position of their pots
            if self.yolo.flag_pair:

                ### --- DEPTH ESTIMATION --- ###

                self.midas.estimate_depth_map(img)

                ### --- DEPTH FUSION --- ###
                
                # lidar scan
                self.depth_fusor.lidar_ranges = np.array(lidar_msg.ranges)
                self.depth_fusor.lidar_range_min = lidar_msg.range_min
                self.depth_fusor.lidar_range_max  = lidar_msg.range_max

                # relative inverse depth map
                self.depth_fusor.depth_map_rel = self.midas.depth_map_rel

                # compute metric depth map
                self.depth_fusor.fuse_depth()

                ### --- POSITION ESTIMATION --- ###

                # bounding boxes of paired pots
                self.depth_fusor.boxes = self.yolo.paired_pots_xyxy

				# position
                self.depth_fusor.get_distance() # TODO show distance in anotations
                self.depth_fusor.get_position()

                ### --- WRITE RESULTS --- ###

                self.result.exitcode = 0
                # TODO vyresit vice objektu
                self.result.pose.position.x = self.depth_fusor.object_positions[0,0]
                self.result.pose.position.y = self.depth_fusor.object_positions[0,1]
                self.result.pose.position.z = self.depth_fusor.object_positions[0,2]
                self.result.pose.orientation.x = 0 # not used
                self.result.pose.orientation.y = 0 # not used
                self.result.pose.orientation.z = 0 # not used
                self.result.pose.orientation.w = 0 # not used

                if not self.DATA_SAVING:
                    break # break from while loop and publish results outside of it

            ### --- ANOTATE IMAGE --- ###

            if self.DO_DEBUGING and self.yolo.flag_any_flower:

                t0 = time.time()

                # anotate original image
                self.yolo.anotate_raw()
                if self.yolo.flag_pair:
                    # anotate pairs with distances, too
                    self.yolo.anotate_pairs_distances(self.depth_fusor.object_distances)

                t1 = time.time()
                # rospy.loginfo(f'Anotation took {np.round(1000*(t1-t0))} ms')

            # publish image
            self.image_pub.publish(self.bridge.cv2_to_imgmsg(self.yolo.img0, "bgr8"))
            if self.DATA_SAVING:
                cv2.imwrite(str(ROOT/'test_files/detections.png'), self.yolo.img0)

            ### --- EXPERIMENTAL LOGGING --- ###

            if self.DATA_SAVING and self.yolo.flag_pair:
                log_path = ROOT / 'logging'

                flag_save = True
                # while True:
                #     flag_save = False
                #     print('\n')
                #     keyboard = input('ENTER to continue, or distances [cm] to save files\n')
                #     if keyboard == '':
                #         break
                #     elif keyboard[0].isdecimal():
                #         flag_save = True
                #         break

                ### save experimental data
                if flag_save == True:

                    now = datetime.now()
                    stamp = now.strftime("%Y_%m_%d_%H_%M_%S_%f")

                    # log_dir / f'{name}_lidar.pkl'       
                    # save distances from terminal (measured with ruler)
                    with open(log_dir / f'{name}_GT_distances.txt', 'r') as f:
                        keyboard = f.read().rstrip()

                    with open(save_dir / f'{name}_GT_distances.txt', 'w') as f:
                        f.write(keyboard)

                    # save lidar beams
                    with open(save_dir / f'{name}_lidar.pkl', 'wb') as f:
                        pickle.dump(lidar_msg, f)

                    # save original image
                    cv2.imwrite(str(save_dir / f'{name}_img_original.png'), img)

                    # save anotated image
                    cv2.imwrite(str(save_dir / f'{name}_img_predictions.png'), self.yolo.img0)

                    # save least squares
                    im = cv2.imread(str(ROOT / 'test_files/regression.png'))
                    cv2.imwrite(str(save_dir / f'{name}_regression.png'), im)

                    # save metric depth map
                    im = cv2.imread(str(ROOT / 'test_files/metric_depth.png'))
                    cv2.imwrite(str(save_dir / f'{name}_metric_depth.png'), im)

                    # save object distances and positions
                    np.savez(save_dir / f'{name}_distances.npz',
                        object_distances = self.depth_fusor.object_distances,
                        object_positions = self.depth_fusor.object_positions,
                        scale = self.depth_fusor.scale,
                        shift = self.depth_fusor.shift
                        )
                    np.savetxt(save_dir / f'{name}_distances.txt',
                        100*self.depth_fusor.object_distances,fmt='%3.2f', delimiter=',')
                    np.savetxt(save_dir / f'{name}_positions.txt',
                        self.depth_fusor.object_positions, delimiter=',')
                    np.savetxt(save_dir / f'{name}_scaleShift.txt',
                        np.array([self.depth_fusor.scale,self.depth_fusor.shift]), delimiter=',')

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

        ### publish results
        self.server.set_succeeded(self.result)
    
    '''
    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()
