import os
import time
import json
import sys
from datetime import datetime
from collections import defaultdict
from ament_index_python.packages import get_package_share_directory

import cv2
import numpy as np
from cv_bridge import CvBridge

import torch
from ultralytics import YOLO

import rclpy
from rclpy.node import Node
from message_filters import ApproximateTimeSynchronizer, Subscriber

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

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.checkpoint_path = None
        self.checkpoint_name = None
        self.function = None
        self.output_boxes = None
        self.output_masks = None
        self.K = self.D = None
        self.map1 = self.map2 = None

        self.fx = self.fy = 1.0
        self.cv_bridge = CvBridge()
        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_publisher()
        self._init_subscriber()

        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('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('output_boxes', True)
        self.output_boxes = self.get_parameter('output_boxes').value

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

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

        self.declare_parameter('color_image_topic', '/camera/camera/color/image_raw')
        self.color_image_topic = self.get_parameter('color_image_topic').value

        self.declare_parameter('depth_image_topic', '/camera/camera/aligned_depth_to_color/image_raw')
        self.depth_image_topic = self.get_parameter('depth_image_topic').value

        self.declare_parameter('camera_info_topic', '/camera/camera/color/camera_info')
        self.camera_info_topic = self.get_parameter('camera_info_topic').value


    def _init_model(self):
        """init model"""
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        try:
            self.model = YOLO(self.checkpoint_path).to(device)
        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}')

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

    def _init_publisher(self):
        """init_publisher"""
        self.pub_pose_list = self.create_publisher(PoseArrayClassAndID, '/pose/cv_detect_pose', 10)

        if self.output_boxes or self.output_masks:
            self.pub_detect_image = self.create_publisher(Image, '/image/detect_image', 10)

    def _init_subscriber(self):
        """init_subscriber"""
        self.sub_camera_info = self.create_subscription(
            CameraInfo,
            self.camera_info_topic,
            self._camera_info_callback,
            10
        )
        time.sleep(1)

        '''sync get color and depth img'''
        self.sub_color_image = Subscriber(self, Image, self.color_image_topic)
        self.sub_depth_image = Subscriber(self, Image, self.depth_image_topic)

        self.sync_subscriber = ApproximateTimeSynchronizer(
            [self.sub_color_image, self.sub_depth_image],
            queue_size=10,
            slop=0.1
        )
        self.sync_subscriber.registerCallback(self._sync_callback)

    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 self.D is not None:
            self.map1, self.map2, self.K = get_map(msg.k, msg.d, self.camera_size)

            if len(self.D) != 0:
                self.destroy_subscription(self.sub_camera_info)
            else:
                self.D = [0, 0, 0, 0, 0, 0, 0, 0]
                self.destroy_subscription(self.sub_camera_info)
        else:
            raise "K and D are not defined"


    def _sync_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_dict = 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_dict is not None:
            pose_list_all = PoseArrayClassAndID()
            for (class_id, class_name), pose_list in pose_dict.items():
                pose_list_all.objects.append(
                    PoseClassAndID(
                        class_name = class_name,
                        class_id = class_id,
                        pose_list = pose_list
                    )
                )
            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 _seg_image(self, rgb_img: np.ndarray, depth_img: np.ndarray):
        """Use segmentation model"""
        pose_dict = defaultdict(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()

        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:
                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]

            intrinsics = [
                int(self.camera_size[0]),
                int(self.camera_size[1]),
                self.K[2] - x_min,
                self.K[5] - y_min,
                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=self.config["command"]["camera_pose"],
                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()}"
            )

            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:
                            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_dict[int(class_ids[i]), labels[class_ids[i]]].append(pose)

                    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_dict


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