#!/usr/bin/env python3

# old node for debugging

import rospy
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

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

class ArucoDetector:

  def __init__(self):
    self.image_pub = rospy.Publisher("detected_aruco_img",Image, queue_size=1)
    self.id_pub = rospy.Publisher("aruco_ID", String, queue_size=1)
    self.bridge = CvBridge()
    self.image_sub = rospy.Subscriber("video_frames", Image, self.callback)

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

    # init aruco utilities                   DICT_ARUCO_ORIGINAL   DICT_5X5_1000
    self.aruco_dict = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_5X5_1000)
    self.marker_length = 0.022 # [m]; moje vytisknute: 0.022 m,  Ondrova simulace: 0.077 (770 px)
    self.axis_length = self.marker_length/2 # [m] for display of a
    self.arucoParams = cv2.aruco.DetectorParameters_create()
    self.arucoParams.minMarkerPerimeterRate = 0.03 # 0.03
    self.arucoParams.maxMarkerPerimeterRate = 4.0 # 4.0

    rospy.loginfo("Path: %s",calib_path)
  

  def callback(self,data):
    cv_image = self.bridge.imgmsg_to_cv2(data)

    markers_img, ids_list = self.detect_aruco(cv_image)

    # publish detected ids
    if ids_list is None:
      self.id_pub.publish(ids_list)
      rospy.loginfo("Detected IDs: None") # debug
    else:
      ids_str = ''.join(str(e) for e in ids_list)
      self.id_pub.publish(ids_str)
      rospy.loginfo("Detected IDs: %s.", ids_str) # debug

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

  def detect_aruco(self,img_rgb):
    img_gray = cv2.cvtColor(img_rgb, cv2.COLOR_BGR2GRAY)

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


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

    # draw markers if any detected
    if ids is not None:
        img_rgb = cv2.aruco.drawDetectedMarkers(img_rgb, corners, ids)
        for rvec, tvec in zip(rvecs, tvecs):
            img_rgb = cv2.aruco.drawAxis(img_rgb, self.cam_mtx, \
              self.dist_coefs, rvec, tvec, self.axis_length)

    return img_rgb, ids

if __name__ == '__main__':
  rospy.init_node('detect_aruco', anonymous=False)
  rospy.loginfo("Initializing ArUco detection node.")
  aruco_det = ArucoDetector()
  rospy.spin()