"""Vision Detect Node"""
import os
import time
import threading
import json
from ament_index_python.packages import get_package_share_directory

import cv2
import open3d as o3d
import numpy as np

import torch
from message_filters import Subscriber, ApproximateTimeSynchronizer

from ultralytics import YOLO

import rclpy
from rclpy.node import Node
from rclpy.task import Future
from cv_bridge import CvBridge

from sensor_msgs.msg import Image, CameraInfo
from geometry_msgs.msg import Pose, Point, Quaternion
from interfaces.msg import PoseClassAndID, PoseArrayClassAndID, ImgMsg
from interfaces.srv import VisionObjectRecognition


from ..utils import calculate_pose_pca, calculate_pose_icp, get_map, \
    distortion_correction, crop_imgs_box_xywh, \
    draw_box, draw_mask, crop_imgs_mask, calculate_grav_width

share_dir = get_package_share_directory('vision_detect')


class DetectNode(Node):
    """Detect Node"""
    def __init__(self, name):
        super().__init__(name)
        self.device = None
        self.checkpoint_path = None
        self.output_boxes = None
        self.output_masks = None

        self.function = None
        self.source = None

        self.server = None
        self.model = None

        self.k = self.d = None
        self.eye_in_left_hand_mat = None
        self.eye_in_right_hand_mat = None
        self.eye_to_hand_mat = None
        self.hand_eye_mat = None

        self.eye_in_left_hand_path = None
        self.eye_in_right_hand_path = None
        self.eye_to_hand_path = None

        self.get_camera_mode = None
        self.detect_mode = None
        self.calculate_mode = None

        self.camera_size = None

        self.map1 = self.map2 = None

        self.calculate_function = calculate_pose_pca
        self.fx = self.fy = 0.5
        self.camera_data = {}
        self.future = Future()
        self.cv_bridge = CvBridge()
        self.lock = threading.Lock()
        self.distance = 1500
        self.color_range = [
            [[0, 120, 70], [10, 255, 255]],
            [[170, 120, 70], [180, 255, 255]]
        ]

        '''init'''
        self._init_param()
        self._init_publisher()
        self._init_subscriber()

        self.get_logger().info("Initialize done")

    def _init_param(self):
        """init parameter"""
        self.declare_parameter('output_boxes', True)
        self.declare_parameter('output_masks', False)
        self.declare_parameter('get_camera_mode', 'service')
        self.declare_parameter('detect_mode', 'detect')
        self.declare_parameter('calculate_mode', 'pca')
        self.declare_parameter('left_hand', 'configs/hand_eye_mat/hand_in_left_hand.json')
        self.declare_parameter('right_hand', 'configs/hand_eye_mat/hand_in_right_hand.json')
        self.declare_parameter('head', 'configs/hand_eye_mat/hand_to_hand.json')

        self.output_boxes = self.get_parameter('output_boxes').value
        self.output_masks = self.get_parameter('output_masks').value
        self.get_camera_mode = self.get_parameter('get_camera_mode').value
        self.detect_mode = self.get_parameter('detect_mode').value
        self.calculate_mode = self.get_parameter("calculate_mode").value

        self._init_json_file()
        if self.get_camera_mode == "Service":
            self.declare_parameter('Service_configs', '{}')
            self.service_name = json.loads(self.get_parameter('Service_configs').value)["service_name"]
            self._init_service()
        elif self.get_camera_mode == "Topic":
            self.declare_parameter('Topic_configs', '{}')
            self.color_image_topic_name = json.loads(self.get_parameter('Topic_configs').value)["color_image_topic_name"]
            self.depth_image_topic_name = json.loads(self.get_parameter('Topic_configs').value)["depth_image_topic_name"]
            self.camera_info_topic_name = json.loads(self.get_parameter('Topic_configs').value)["camera_info_topic_name"]
            self.position = json.loads(self.get_parameter('Topic_configs').value)["position"]
            if self.position == 'left':
                self.hand_eye_mat = self.eye_in_left_hand_mat
            elif self.position == 'right':
                self.hand_eye_mat = self.eye_in_right_hand_mat
            else:
                self.hand_eye_mat = self.eye_to_hand_mat
        else:
            self.declare_parameter('Service_configs', '{}')
            self.service_name = json.loads(self.get_parameter('Service_configs').value)["service_name"]

        if self.detect_mode == 'Detect':
            self.declare_parameter('Detect_configs', '{}')
            self.function = self._seg_image
            self.confidence = json.loads(self.get_parameter('Detect_configs').value)["confidence"]
            self.classes = json.loads(self.get_parameter('Detect_configs').value)["classes"]
            if not self.classes:
                self.classes = None
            self.checkpoint_path = os.path.join(share_dir, json.loads(self.get_parameter('Detect_configs').value)["checkpoint_path"])
            self._init_model()
        elif self.detect_mode == 'Color':
            self.declare_parameter('Color_configs', '{}')
            self.function = self._seg_color
            self.color_range = json.loads(self.get_parameter('Color_configs').value)["color_range"]
            self.color_range = [[np.array(lower), np.array(upper)] for lower, upper in self.color_range]
        elif self.detect_mode == 'Crossboard':
            self.declare_parameter('Crossboard_configs', '{}')
            self.function = self._seg_crossboard
            self.pattern_size = json.loads(self.get_parameter('Crossboard_configs').value)["pattern_size"]
        else:
            self.get_logger().warning("Unknown mode, use detect")
            self.declare_parameter('Detect_configs', '{}')
            self.function = self._seg_image
            self.confidence = json.loads(self.get_parameter('Detect_configs').value)["confidence"]
            self.classes = json.loads(self.get_parameter('Detect_configs').value)["classes"]
            if not self.classes:
                self.classes = None
            self.checkpoint_path = json.loads(self.get_parameter('Detect_configs').value)["checkpoint_path"]
            self._init_model()

        if self.calculate_mode == "PCA":
            self.declare_parameter('PCA_configs', '{}')
            self.configs = json.loads(self.get_parameter('PCA_configs').value)
            self.calculate_function = calculate_pose_pca
        elif self.calculate_mode == "ICP":
            self.declare_parameter('ICA_configs', '{}')
            self.configs = json.loads(self.get_parameter('ICA_configs').value)
            self.source = o3d.io.read_point_cloud(os.path.join(share_dir, self.configs['complete_model_path']))
            self.calculate_function = calculate_pose_icp
        else:
            self.get_logger().warning("Unknown calculate_mode, use PCA")
            self.declare_parameter('PCA_configs', '{}')
            self.configs = json.loads(self.get_parameter('PCA_configs').value)
            self.calculate_function = calculate_pose_pca

        self.get_logger().info("Initialize parameters done")

    def _init_json_file(self):
        self.eye_in_left_hand_path = os.path.join(share_dir, self.get_parameter('left_hand').value)
        self.eye_in_right_hand_path = os.path.join(share_dir, self.get_parameter('right_hand').value)
        self.eye_to_hand_path = os.path.join(share_dir, self.get_parameter('head').value)

        if not os.path.exists(self.eye_in_left_hand_path):
            self.get_logger().warning(f"{self.eye_in_left_hand_path} not found, use E(4, 4)")
            self.eye_in_left_hand_mat = np.eye(4)
        else:
            with open(self.eye_in_left_hand_path, "r", encoding="utf-8") as f:
                self.eye_in_left_hand_mat = np.array(json.load(f)["T"])
                if self.eye_in_left_hand_mat.shape != (4, 4):
                    self.get_logger().warning("The shape of eye_in_right_hand_mat is wrong, use E(4, 4)")
                    self.eye_in_left_hand_mat = np.eye(4)

        if not os.path.exists(self.eye_in_right_hand_path):
            self.get_logger().warning(f"{self.eye_in_right_hand_path} not found, use E(4, 4)")
            self.eye_in_right_hand_mat = np.eye(4)
        else:
            with open(self.eye_in_right_hand_path, "r", encoding="utf-8") as f:
                self.eye_in_right_hand_mat = np.array(json.load(f)["T"])
                if self.eye_in_right_hand_mat.shape != (4, 4):
                    self.get_logger().warning("The shape of eye_in_right_hand_mat is wrong, use E(4, 4)")
                    self.eye_in_right_hand_mat = np.eye(4)

        if not os.path.exists(self.eye_to_hand_path):
            self.get_logger().warning(f"{self.eye_to_hand_path} not found, use E(4, 4)")
            self.eye_to_hand_mat = np.eye(4)
        else:
            with open(self.eye_to_hand_path, "r", encoding="utf-8") as f:
                self.eye_to_hand_mat = np.array(json.load(f)["T"])
                if self.eye_to_hand_mat.shape != (4, 4):
                    self.get_logger().warning("The shape of eye_to_hand_mat is wrong, use E(4, 4)")
                    self.eye_to_hand_mat = np.eye(4)

        self.get_logger().info("Initialize read json file done")

    def _init_model(self):
        """init model"""
        device_model = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        try:
            self.model = YOLO(self.checkpoint_path).to(device_model)
        except Exception as e:
            self.get_logger().error(f'Failed to load YOLO model: {e}')
            raise 
        self.get_logger().info(f'Loading checkpoint from: {self.checkpoint_path}')

        if self.checkpoint_path.endswith('-seg.pt'):
            self.function = self._seg_image
        else:
            self.function = None
            self.get_logger().fatal(f'Unknown checkpoint: {self.checkpoint_path}')
            raise ValueError(f'Unknown checkpoint: {self.checkpoint_path}')
        self.get_logger().info("Initialize model done")

    def _init_publisher(self):
        """init publisher"""
        if self.output_boxes or self.output_masks:
            self.pub_detect_image = self.create_publisher(Image, '/image/detect_image', 10)

        if self.get_camera_mode == "Topic":
            self.pub_pose_list = self.create_publisher(PoseArrayClassAndID, '/pose/cv_detect_pose', 10)

        self.get_logger().info("Initialize publisher done")

    def _init_service(self):
        """init service server"""
        self.server = self.create_service(
            VisionObjectRecognition,
            self.service_name,
            self._service_callback
        )
        self.get_logger().info("Initialize service done")

    def _init_subscriber(self):
        """init subscriber"""
        if self.get_camera_mode == "Service":
            self.sub_img = self.create_subscription(
                ImgMsg,
                "/img_msg",
                self._service_sub_callback,
                10
            )
        elif self.get_camera_mode == "Topic":
            self.sub_camera_info = self.create_subscription(
                CameraInfo,
                self.camera_info_topic_name,
                self._camera_info_callback,
                10
            )

            self.get_logger().info("Waiting for camera info...")
            rclpy.spin_until_future_complete(self, self.future)
            self.get_logger().info("Camera info received.")

            # sync get color and depth img
            self.sub_color_image = Subscriber(self, Image, self.color_image_topic_name)
            self.sub_depth_image = Subscriber(self, Image, self.depth_image_topic_name)

            self.sync_subscriber = ApproximateTimeSynchronizer(
                [self.sub_color_image, self.sub_depth_image],
                queue_size=10,
                slop=0.1
            )
            self.sync_subscriber.registerCallback(self._sync_sub_callback)
        else:
            self.get_logger().warning("get_camera_mode is wrong, use service")
            self.sub_img = self.create_subscription(
                ImgMsg,
                "/img_msg",
                self._service_sub_callback,
                10
            )

        self.get_logger().info("Initialize subscriber done")

    def _camera_info_callback(self, msg: CameraInfo):
        """Get camera info"""
        self.k = msg.k
        self.d = msg.d

        self.camera_size = [msg.width, msg.height]

        if self.k is not None and len(self.k) > 0 and self.d is not None and len(self.d) > 0:
            self.map1, self.map2, self.k = get_map(msg.k, msg.d, self.camera_size)
            if not self.future.done(): 
                self.future.set_result(True)
            self.destroy_subscription(self.sub_camera_info)
        else:
            self.get_logger().warning("K and d are not defined")

    def _service_sub_callback(self, msgs):
        """同步回调函数"""
        with self.lock:
            self.camera_data[msgs.position] = [
                msgs.image_color,
                msgs.image_depth,
                msgs.karr,
                msgs.darr
            ]

    def _sync_sub_callback(self, color_img_ros, depth_img_ros):
        """同步回调函数"""
        color_img_cv = self.cv_bridge.imgmsg_to_cv2(color_img_ros, "bgr8")
        depth_img_cv = self.cv_bridge.imgmsg_to_cv2(depth_img_ros, '16UC1')

        color_img_cv, depth_img_cv = distortion_correction(color_img_cv, depth_img_cv, self.map1, self.map2)

        img, pose_list = self.function(color_img_cv, depth_img_cv)

        # masks为空，结束这一帧
        if img is None:
            img = self.cv_bridge.cv2_to_imgmsg(color_img_cv, "bgr8")

        if self.output_boxes or self.output_masks:
            self.pub_detect_image.publish(img)

        if pose_list:
            pose_list_all = PoseArrayClassAndID()
            for item in pose_list:
                pose_list_all.objects.append(
                    PoseClassAndID(
                        class_name = item["class_name"],
                        class_id = item["class_id"],
                        pose = item["pose"],
                        grab_width = item["grab_width"]
                    )
                )
            pose_list_all.header.stamp = self.get_clock().now().to_msg()
            pose_list_all.header.frame_id = "pose_list"
            self.pub_pose_list.publish(pose_list_all)

    def _service_callback(self, request, response):
        response.header.stamp = self.get_clock().now().to_msg()
        response.header.frame_id = "camera_detect"

        with self.lock:
            if request.camera_position in self.camera_data:
                color_img_ros, depth_img_ros, self.k, d = self.camera_data[request.camera_position]
            else:
                response.success = False
                response.info = f"{request.camera_position} Camera data is empty or name is wrong"
                response.objects = []

                return response

        if request.camera_position == 'left':
            self.hand_eye_mat = self.eye_in_left_hand_mat
        elif request.camera_position == 'right':
            self.hand_eye_mat = self.eye_in_right_hand_mat
        else:
            self.hand_eye_mat = self.eye_to_hand_mat

        self.camera_size = [color_img_ros.width, color_img_ros.height]

        color_img_cv = self.cv_bridge.imgmsg_to_cv2(color_img_ros, "bgr8")
        depth_img_cv = self.cv_bridge.imgmsg_to_cv2(depth_img_ros, '16UC1')

        map1, map2, self.k = get_map(self.k, d, self.camera_size)
        color_img_cv, depth_img_cv = distortion_correction(color_img_cv, depth_img_cv, map1, map2)

        img, pose_list = self.function(color_img_cv, depth_img_cv)

        # masks为空，结束这一帧
        if self.output_boxes or self.output_masks:
            if img is None:
                img = color_img_ros
            self.pub_detect_image.publish(img)

        if pose_list:
            response.info = "Success get pose"
            response.success = True
            for item in pose_list:
                response.objects.append(
                    PoseClassAndID(
                        class_name = item["class_name"],
                        class_id = item["class_id"],
                        pose = item["pose"],
                        grab_width = item["grab_width"]
                    )
                )
        else:
            response.info = "pose dict is empty"
            response.success = False
            response.objects = []

        return response

    def _seg_image(self, rgb_img: np.ndarray, depth_img: np.ndarray):
        """Use segmentation model"""
        pose_list = []

        # Get Predict Results
        time1 = time.time()
        results = self.model(rgb_img, retina_masks=True, conf=self.confidence, classes=self.classes)
        time2 = time.time()
        result = results[0]

        # Get masks
        if result.masks is None or len(result.masks) == 0:
            return None, None
        masks = result.masks.data.cpu().numpy()

        # Get boxes
        boxes = result.boxes.xywh.cpu().numpy()
        class_ids = result.boxes.cls.cpu().numpy()
        labels = result.names

        time3 = time.time()

        for i, (mask, box) in enumerate(zip(masks, boxes)):
            imgs_crop, (x_min, y_min) = crop_imgs_box_xywh([depth_img, mask], box, True)
            depth_crop, mask_crop = imgs_crop

            if depth_crop is None:
                continue

            intrinsics = o3d.camera.PinholeCameraIntrinsic(
                int(self.camera_size[0]),
                int(self.camera_size[1]),
                self.k[0],
                self.k[4],
                self.k[2] - x_min,
                self.k[5] - y_min
            )

            x, y, z, rw, rx, ry, rz = self.calculate_function(
                mask_crop,
                depth_crop,
                intrinsics,
                self.source,
                self.hand_eye_mat,
                self.configs
            )

            grab_width = calculate_grav_width(mask, self.k, z)
            z = z + grab_width * 0.45

            if (x, y, z) != (0.0, 0.0, 0.0):
                pose = Pose()
                pose.position = Point(x=x, y=y, z=z)
                pose.orientation = Quaternion(w=rw, x=rx, y=ry, z=rz)
                pose_list.append(
                    {
                        "class_id": int(class_ids[i]),
                        "class_name": labels[class_ids[i]],
                        "pose": pose,
                        "grap_width": grab_width
                    }
                )

        time4 = time.time()

        self.get_logger().info('start')
        self.get_logger().info(f'{(time2 - time1) * 1000} ms, model predict')
        self.get_logger().info(f'{(time3 - time2) * 1000} ms, get mask and some param')
        self.get_logger().info(f'{(time4 - time3) * 1000} ms, calculate all mask PCA')
        self.get_logger().info(f'{(time4 - time1) * 1000} ms, completing a picture entire process')
        self.get_logger().info('end')

        # mask_img and box_img is or not output
        if self.output_boxes and not self.output_masks:
            draw_box(rgb_img, result)
            return self.cv_bridge.cv2_to_imgmsg(rgb_img, "bgr8"), pose_list
        elif self.output_boxes and self.output_masks:
            draw_box(rgb_img, result)
            draw_mask(rgb_img, result)
            return self.cv_bridge.cv2_to_imgmsg(rgb_img, "bgr8"), pose_list
        elif not self.output_boxes and self.output_masks:
            draw_mask(rgb_img, result)
            return self.cv_bridge.cv2_to_imgmsg(rgb_img, "bgr8"), pose_list
        else:
            return None, pose_list

    def _seg_color(self, rgb_img: np.ndarray, depth_img: np.ndarray):
        """Use segmentation model"""
        pose_list = []

        hsv_img = cv2.cvtColor(rgb_img, cv2.COLOR_BGR2HSV)

        depth_filter_mask = np.zeros_like(depth_img, dtype=np.uint8)
        depth_filter_mask[(depth_img > 0) & (depth_img < self.distance)] = 1

        hsv_img[depth_filter_mask == 0] = 0

        mask_part_list = [cv2.inRange(hsv_img, color[0], color[1]) for color in self.color_range]
        mask = sum(mask_part_list[1:], mask_part_list[0])
        mask = mask // 255

        rgb_img[mask > 0] = rgb_img[mask > 0] * 0.5 + np.array([255, 0, 0]) * 0.5

        imgs_crop, mask_crop, (x_min, y_min) = crop_imgs_mask([depth_img], mask, True)
        depth_crop = imgs_crop[0]

        if depth_crop is None:
            self.get_logger().warning("depth_crop is None")
            return None, None

        intrinsics = o3d.camera.PinholeCameraIntrinsic(
            int(self.camera_size[0]),
            int(self.camera_size[1]),
            self.k[0],
            self.k[4],
            self.k[2] - x_min,
            self.k[5] - y_min
        )

        x, y, z, rw, rx, ry, rz = self.calculate_function(
            mask_crop,
            depth_crop,
            intrinsics,
            self.source,
            self.hand_eye_mat,
            self.configs
        )

        if (x, y, z) != (0.0, 0.0, 0.0):
            pose = Pose()
            pose.position = Point(x=x, y=y, z=z)
            pose.orientation = Quaternion(w=rw, x=rx, y=ry, z=rz)
            pose_list.append(
                {
                    "class_id": int(98),
                    "class_name": "red",
                    "pose": pose,
                }
            )

        return self.cv_bridge.cv2_to_imgmsg(rgb_img, "bgr8"), pose_list

    def _seg_crossboard(self, rgb_img, depth_img):
        pose_list = []
        rgb_img_gray = cv2.cvtColor(rgb_img, cv2.COLOR_BGR2GRAY)
        ret, corners = cv2.findChessboardCorners(rgb_img_gray, self.pattern_size, cv2.CALIB_CB_FAST_CHECK)
        if ret:
            # 角点亚像素精确化（提高标定精度）
            criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
            corners_subpix = cv2.cornerSubPix(rgb_img_gray, corners, (11, 11), (-1, -1), criteria)

            corners_subpix = corners_subpix.reshape(self.pattern_size[1], self.pattern_size[0], 2)
            mask = np.zeros(rgb_img_gray.shape, dtype=np.uint8)

            for i in range(0, self.pattern_size[1] - 1):
                for j in range(0, self.pattern_size[0] - 1):
                    pts = np.array([
                        corners_subpix[i, j],
                        corners_subpix[i, j + 1],
                        corners_subpix[i + 1, j + 1],
                        corners_subpix[i + 1, j]
                    ], dtype=np.int32)
                    cv2.fillConvexPoly(mask, pts, 1)

            rgb_img[mask > 0] = rgb_img[mask > 0] * 0.5 + np.array([0, 0, 255]) * 0.5

            img_crop, mask_crop, (x_min, y_min) = crop_imgs_mask([depth_img], mask)
            depth_crop = img_crop[0]

            if depth_crop is None:
                return None, None

            intrinsics = o3d.camera.PinholeCameraIntrinsic(
                int(self.camera_size[0]),
                int(self.camera_size[1]),
                self.k[0],
                self.k[4],
                self.k[2] - x_min,
                self.k[5] - y_min
            )

            x, y, z, rw, rx, ry, rz = self.calculate_function(
                mask_crop,
                depth_crop,
                intrinsics,
                self.source,
                self.hand_eye_mat,
                self.configs
            )

            self.get_logger().info(f"{x}, {y}, {z}, {rw}, {rx}, {ry}, {rz}")

            if (x, y, z) != (0.0, 0.0, 0.0):
                pose = Pose()
                pose.position = Point(x=x, y=y, z=z)
                pose.orientation = Quaternion(w=rw, x=rx, y=ry, z=rz)
                pose_list.append(
                    {
                        "class_id": int(99),
                        "class_name": 'crossboard',
                        "pose": pose,
                    }
                )

            cv2.putText(
                rgb_img,
                f'cs: x: {x:.3f}, y: {y:.3f}, z: {z:.3f}',
                (0, 0 + 30),cv2.FONT_HERSHEY_SIMPLEX,
                1,
                (255, 255, 0),
                2
            )
            cv2.putText(
                rgb_img,
                f'quat: rw: {rw:.3f}, rx: {rx:.3f}, ry: {ry:.3f}, rz: {rz:.3f}',
                (0, 0 + 70), cv2.FONT_HERSHEY_SIMPLEX,
                1,
                (255, 255, 0),
                2
            )
            return self.cv_bridge.cv2_to_imgmsg(rgb_img, "bgr8"), pose_list
            
        return self.cv_bridge.cv2_to_imgmsg(rgb_img, "bgr8"), None


if __name__ == '__main__':
    rclpy.init(args=None)
    node = DetectNode('detect')
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()
