import os
import time
import threading
import json
import sys
import ast
from datetime import datetime
from ament_index_python.packages import get_package_share_directory

import cv2
import numpy as np
import transforms3d as tfs
from cv_bridge import CvBridge

import torch
from ultralytics import YOLO

import rclpy
from rclpy.node import Node

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

sys.path.insert(0, "/home/nvidia/aidk/AIDK/release/lib_py/")
import flexivaidk as aidk

share_dir = get_package_share_directory('vision_detect')


def get_map(K, D, camera_size):
    h, w = camera_size[::-1]
    K = np.array(K).reshape(3, 3)
    D = np.array(D)
    new_K, _ = cv2.getOptimalNewCameraMatrix(K, D, (w, h), 1, (w, h))
    map1, map2 = cv2.initUndistortRectifyMap(K, D, None, new_K, (w, h), cv2.CV_32FC1)

    return map1, map2, new_K.flatten()


def distortion_correction(color_image, depth_image, map1, map2):
    """畸变矫正"""
    undistorted_color = cv2.remap(color_image, map1, map2, cv2.INTER_LINEAR)
    undistorted_color = undistorted_color.astype(color_image.dtype)

    undistorted_depth = cv2.remap(depth_image, map1, map2, cv2.INTER_NEAREST)
    undistorted_depth = undistorted_depth.astype(depth_image.dtype)

    return undistorted_color, undistorted_depth


def crop_mask_bbox(rgb_img, depth_img, mask, box):
    """
    输入:
        depth_img: H x W
        mask: H x W (0/1 或 bool)
    输出:
        depth_crop, mask_crop
    """
    high, width = depth_img.shape
    x_center, y_center, w, h = box[:4]

    x_min, x_max = int(round(x_center - w/2)), int(round(x_center + w/2))
    y_min, y_max = int(round(y_center - h/2)), int(round(y_center + h/2))

    rgb_crop = rgb_img[max(0, y_min):min(y_max, high) + 1, max(0, x_min):min(x_max, width) + 1]
    depth_crop = depth_img[max(0, y_min):min(y_max, high) + 1, max(0, x_min):min(x_max, width) + 1]
    mask_crop = mask[max(0, y_min):min(y_max, high) + 1, max(0, x_min):min(x_max, width) + 1]

    return rgb_crop, depth_crop, mask_crop, (max(0, x_min), max(0, y_min))


class DetectNode(Node):
    def __init__(self, name):
        super().__init__(name)
        self.device = None
        self.checkpoint_path = None
        self.checkpoint_name = None
        self.output_boxes = None
        self.output_masks = None
        self.function = None
        self.K = None
        self.eye_in_hand_mat = None
        self.eye_to_hand_mat = None

        self.fx = self.fy = 1.0
        self.camera_data = {}
        self.cv_bridge = CvBridge()
        self.lock = threading.Lock()

        self.aidk_client = aidk.AIDKClient('127.0.0.1', 10)
        while not self.aidk_client.is_ready():
            time.sleep(0.5)
        
        '''init'''
        self._init_param()
        self._init_model()
        self._init_config()
        self._init_json_file()
        self._init_publisher()
        self._init_subscriber()
        self._init_service()

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

    def _init_param(self):
        """init parameter"""
        self.declare_parameter('checkpoint_name', 'yolo11s-seg.pt')
        self.checkpoint_name = self.get_parameter('checkpoint_name').value
        self.checkpoint_path = os.path.join(share_dir, 'checkpoints', self.checkpoint_name)

        self.declare_parameter('output_boxes', True)
        self.output_boxes = self.get_parameter('output_boxes').value

        self.declare_parameter('output_masks', False)
        self.output_masks = self.get_parameter('output_masks').value

        self.declare_parameter('config_name', 'default_config.json')
        self.config_name = self.get_parameter('config_name').value
        self.config_dir = os.path.join(share_dir, 'configs/flexiv_configs', self.config_name)

        self.declare_parameter('set_confidence', 0.25)
        self.set_confidence = self.get_parameter('set_confidence').value

        self.declare_parameter('classes', 'None')
        self.classes = ast.literal_eval(self.get_parameter('classes').value)

        self.declare_parameter('eye_in_left_hand_path', "configs/hand_eye_mat/eye_in_left_hand.json")
        self.declare_parameter('eye_in_right_hand_path', "configs/hand_eye_mat/eye_in_right_hand.json")
        self.declare_parameter('eye_to_hand_path', "configs/hand_eye_mat/eye_to_hand.json")

        self.get_logger().info("Init param 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_name.endswith('-seg.pt'):
            self.function = self._seg_image
        else:
            self.function = None
            self.get_logger().error(f'Unknown checkpoint: {self.checkpoint_name}')

        self.get_logger().info("Init 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)
        
        self.get_logger().info("Init publisher done")

    def _init_service(self):
        """init service server"""
        self.server = self.create_service(
            VisionObjectRecognition, 
            "/vision_object_recognition", 
            self._service_callback
        )

        self.get_logger().info("Init service done")

    def _init_subscriber(self):
        """init subscriber"""
        self.sub_img = self.create_subscription(
            ImgMsg,
            "/img_msg",
            self._sub_callback, 
            10
        )

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

    def _init_config(self):
        """init config"""
        with open(self.config_dir, "r") as f:
            self.config = json.load(f)
        
        self.get_logger().info("Init config done")

    def _init_json_file(self):
        self.eye_in_left_hand_path = os.path.join(share_dir, self.get_parameter('eye_in_left_hand_path').value)
        self.eye_in_right_hand_path = os.path.join(share_dir, self.get_parameter('eye_in_right_hand_path').value)
        self.eye_to_hand_path = os.path.join(share_dir, self.get_parameter('eye_to_hand_path').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") 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)
        x, y, z = self.eye_in_left_hand_mat[0:3, 3:4].flatten()
        rw, rx, ry, rz = tfs.quaternions.mat2quat(self.eye_in_left_hand_mat[0:3, 0:3])
        self.eye_in_left_hand_mat = [x, y, z, rw, rx, ry, rz]

        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") 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)
        x, y, z = self.eye_in_right_hand_mat[0:3, 3:4].flatten()
        rw, rx, ry, rz = tfs.quaternions.mat2quat(self.eye_in_right_hand_mat[0:3, 0:3])
        self.eye_in_right_hand_mat = [x, y, z, rw, rx, ry, rz]

        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") 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)
        x, y, z = self.eye_to_hand_mat[0:3, 3:4].flatten()
        rw, rx, ry, rz = tfs.quaternions.mat2quat(self.eye_to_hand_mat[0:3, 0:3])
        self.eye_to_hand_mat = [x, y, z, rw, rx, ry, rz]

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

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

    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 = []
                self.get_logger().warning(f"{request.camera_position} Camera data is empty or name is wrong")
                return response
            
        if request.camera_position == 'left':
            hand_eye_mat = self.eye_in_left_hand_mat
        elif request.camera_position == 'right':
            hand_eye_mat = self.eye_in_right_hand_mat
        else:
            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, hand_eye_mat)

        """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 = []
        self.get_logger().info("Detection Completed")
        return response

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

        '''Get Predict Results'''
        time1 = time.time()
        results = self.model(rgb_img, retina_masks=True, classes=[39], conf=self.set_confidence)
        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()
        orig_shape = result.masks.orig_shape

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

        time3 = time.time()

        rgb_bytes = cv2.imencode('.png', rgb_img)[1]
        depth_bytes = cv2.imencode('.png', depth_img)[1]

        for i, (mask, box) in enumerate(zip(masks, boxes)):
            mask = cv2.resize(mask.astype(np.uint8), orig_shape[::-1], interpolation=cv2.INTER_NEAREST)
            # rgb_crop, depth_crop, mask_crop, (x_min, y_min) = crop_mask_bbox(rgb_img, depth_img, mask, box)

            # if depth_crop is None:
            #     self.get_logger().error("depth_crop is None")
            #     continue

            # depth_img_crop_mask = np.zeros_like(depth_crop)
            # depth_img_crop_mask[mask_crop > 0] = depth_crop[mask_crop > 0]

            # print(rgb_crop.shape)
            # print(rgb_crop.dtype)


            # rgb_bytes = cv2.imencode('.png', rgb_crop)[1]
            # depth_bytes = cv2.imencode('.png', depth_img_crop_mask)[1]

            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            mask_contours = contours[0].reshape(1, -1, 2)
            res = self.aidk_client.set_direct_setting_variables({"one_mask": f"{mask_contours}"})

            intrinsics = [
                int(self.camera_size[0]),
                int(self.camera_size[1]),
                # self.K[2] - x_min,
                # self.K[5] - y_min,
                self.K[2],
                self.K[5],
                self.K[0],
                self.K[4]
            ]

            state = self.aidk_client.detect_with_image(
                obj_name=self.config["command"]["obj_name"],
                camera_id=self.config["command"]["camera_id"],
                coordinate_id=self.config["command"]["coordinate_id"],
                camera_pose=hand_eye_mat,
                camera_intrinsic=intrinsics,
                rgb_input=aidk.ImageBuffer(rgb_bytes),
                depth_input=aidk.ImageBuffer(depth_bytes),
                custom=self.config["command"]["custom"],
            )

            self.get_logger().info(f"state: {state}")
            self.get_logger().info(
                f"current detected object names: {self.aidk_client.get_detected_obj_names()}, current detected object nums: {self.aidk_client.get_detected_obj_nums()}"
            )

            self.get_logger().info(
                "detected time stamp: {}".format(
                    datetime.fromtimestamp(self.aidk_client.get_detected_time())
                )
            )

            parse_state, result_list_grab_width = self.aidk_client.parse_result(self.config["command"]["obj_name"], "double_value", -1)
            if not parse_state:
                self.get_logger().error("Parse result error!!!")
                continue

            parse_state, result_list_obj_pose = self.aidk_client.parse_result(self.config["command"]["obj_name"], "obj_pose", -1)
            if not parse_state:
                self.get_logger().error("Parse result error!!!")
                continue

            for result_pose, result_width in zip(result_list_obj_pose, result_list_grab_width):
                self.get_logger().info(f'"double_value": {getattr(result_width, "double_value")}')
                for vec in result_pose.vect:
                    self.get_logger().info(f"vec: {vec}")
                    x, y, z, rw, rx, ry, rz = vec

                    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,
                            "grab_width": getattr(result_width, "double_value")
                        }
                    )

            # for key in self.config["keys"]:
            #     parse_state, result_list = self.aidk_client.parse_result(self.config["command"]["obj_name"], key, -1)
            #     self.get_logger().info(
            #         "detected time stamp: {}".format(
            #             datetime.fromtimestamp(self.aidk_client.get_detected_time())
            #         )
            #     )
            #     if not parse_state:
            #         self.get_logger().error("Parse result error!!!")
            #         continue
            #     else:
            #         if key in ["bbox", "keypoints", "positions", "obj_pose"]:
            #             for result in result_list:
            #                 self.get_logger().info(f'"double_value in result": {getattr(result, "double_value")}')
            #                 for vec in result.vect:
            #                     self.get_logger().info(f"vec: {vec}")
            #                     x, y, z, rw, rx, ry, rz = vec
            #
            #                     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,
            #                             "grab_width": getattr(result, "double_value")
            #                         }
            #                     )
            #                     self.get_logger().info(f"pose: position({x}, {y}, {z}), orientation({rw}, {rx}, {ry}, {rz})")
            #                     self.get_logger().info(f"class_id: {int(class_ids[i])}, class_name: {labels[class_ids[i]]}, grab_width: {getattr(result, 'double_value')}")
            #
            #         elif key in ["valid", "double_value", "int_value", "name"]:
            #             for result in result_list:
            #                 self.get_logger().info(f"{key}: {getattr(result, key)}")

        time4 = time.time()

        self.get_logger().info(f'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(f'end')

        return None, pose_list


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


if __name__ == '__main__':
    main()
