#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image, CameraInfo
from cv_bridge import CvBridge
import cv2
import numpy as np
from ultralytics import YOLO
import open3d as o3d
import message_filters
from PIL import Image as PILImage, ImageDraw, ImageFont

class YoloRealSenseFusionNode(Node):
    def __init__(self):
        super().__init__('yolo_realsense_fusion_node')
        
        # 1. 加载 YOLO 模型
        try:
            self.model = YOLO("yolo11n.pt") 
        except Exception as e:
            self.get_logger().error(f'Failed to load model: {e}')
            raise e

        self.bridge = CvBridge()

        # 2. 订阅话题 (确保 realsense-ros 启动时设置了 align_depth:=true)
        self.image_sub = message_filters.Subscriber(self, Image, '/camera/camera/color/image_raw')
        self.depth_sub = message_filters.Subscriber(self, Image, '/camera/camera/depth/image_rect_raw')
        # 订阅相机内参
        self.info_sub = self.create_subscription(CameraInfo, '/camera/camera/depth/camera_info', self.info_callback, 10)

        # 时间同步器：确保 RGB 和深度帧匹配
        self.ts = message_filters.ApproximateTimeSynchronizer(
            [self.image_sub, self.depth_sub], 
            queue_size=10, 
            slop=0.1
        )
        self.ts.registerCallback(self.callback)

        # 发布话题
        self.image_pub = self.create_publisher(Image, 'detect', 10)

        self.intrinsics = None
        self.get_logger().info('YOLO RealSense 融合节点已启动')
        
        # 初始化中文字体
        try:
            self.font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 20)
        except:
            try:
                self.font = ImageFont.truetype("/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc", 20)
            except:
                self.font = ImageFont.load_default()
                self.get_logger().warn('无法加载中文字体，使用默认字体')

    def calculate_object_dimensions(self, pcd):
        """
        计算点云的尺寸（长、宽、高）
        使用轴对齐边界框（AABB）计算
        """
        if pcd.is_empty() or len(pcd.points) == 0:
            return None, None, None
        
        try:
            # 获取轴对齐边界框
            bbox = pcd.get_axis_aligned_bounding_box()
            
            # 获取边界框的尺寸
            size = bbox.get_extent()
            
            # size[0]: X方向宽度（左右）
            # size[1]: Y方向高度（上下）
            # size[2]: Z方向深度（前后）
            width = size[0]
            height = size[1]
            depth = size[2]
            
            return width, height, depth
        except Exception as e:
            self.get_logger().warn(f'计算尺寸失败: {e}')
            return None, None, None

    def info_callback(self, msg):
        # 初始化 Open3D 相机内参
        if self.intrinsics is None:
            self.intrinsics = o3d.camera.PinholeCameraIntrinsic(
                msg.width, msg.height, msg.k[0], msg.k[4], msg.k[2], msg.k[5]
            )
            self.get_logger().info('相机内参已初始化')

    def callback(self, color_msg, depth_msg):
        if self.intrinsics is None:
            return

        try:
            # 转换 ROS 图像为 OpenCV 格式
            color_image = self.bridge.imgmsg_to_cv2(color_msg, desired_encoding='bgr8')
            depth_image = self.bridge.imgmsg_to_cv2(depth_msg, desired_encoding='16UC1')

            # YOLO 推理
            results = self.model(color_image, conf=0.5, verbose=False)
            boxes = results[0].boxes.xyxy.cpu().numpy()
            classes = results[0].boxes.cls.cpu().numpy()
            names = results[0].names

            annotated_image = color_image.copy()

            for box, cls in zip(boxes, classes):
                x1, y1, x2, y2 = map(int, box)
                obj_name = names[int(cls)]

                # 边界框有效性检查
                img_height, img_width = color_image.shape[:2]
                # 确保边界框在图像范围内
                x1 = max(0, min(x1, img_width - 1))
                y1 = max(0, min(y1, img_height - 1))
                x2 = max(0, min(x2, img_width - 1))
                y2 = max(0, min(y2, img_height - 1))
                
                # 确保边界框有有效尺寸
                if x2 <= x1 or y2 <= y1:
                    self.get_logger().warn(f"边界框尺寸无效：({x1},{y1}) -> ({x2},{y2})，跳过此目标")
                    continue

                # 提取目标区域的深度和彩色图
                # 使用 np.ascontiguousarray 确保内存连续，解决 Open3D 初始化错误
                roi_depth = np.ascontiguousarray(depth_image[y1:y2, x1:x2])
                roi_color = np.ascontiguousarray(color_image[y1:y2, x1:x2])

                # 检查并调整ROI尺寸不匹配问题
                if roi_color.shape[:2] != roi_depth.shape[:2]:
                    self.get_logger().debug(f"ROI尺寸不匹配：彩色{roi_color.shape[:2]}，深度{roi_depth.shape[:2]}")
                    # 检查深度ROI的有效尺寸，避免resize错误
                    if roi_depth.shape[0] > 0 and roi_depth.shape[1] > 0:
                        # 将彩色ROI调整到深度ROI的尺寸
                        roi_color = cv2.resize(roi_color, (roi_depth.shape[1], roi_depth.shape[0]), interpolation=cv2.INTER_LINEAR)
                    else:
                        self.get_logger().warn(f"深度ROI尺寸无效：{roi_depth.shape[:2]}，跳过此目标")
                        continue

                # 默认绘制框和名称
                cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 165, 255), 2)
                label_line1 = f"{obj_name}"
                label_line2 = None

                if np.sum(roi_depth > 0) >= 50:
                    
                    # 构造局部点云
                    o3d_color = o3d.geometry.Image(cv2.cvtColor(roi_color, cv2.COLOR_BGR2RGB))
                    o3d_depth = o3d.geometry.Image(roi_depth)

                    # 调整内参的主点坐标以适应 ROI 切片
                    local_intrinsics = o3d.camera.PinholeCameraIntrinsic(
                        roi_depth.shape[1], roi_depth.shape[0],
                        self.intrinsics.intrinsic_matrix[0, 0],
                        self.intrinsics.intrinsic_matrix[1, 1],
                        self.intrinsics.intrinsic_matrix[0, 2] - x1,
                        self.intrinsics.intrinsic_matrix[1, 2] - y1
                    )

                    # 生成 RGBD 图像并转换为点云
                    rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(
                        o3d_color, o3d_depth, depth_scale=1000.0, depth_trunc=3.0, convert_rgb_to_intensity=False
                    )
                    pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd, local_intrinsics)

                    if not pcd.is_empty():
                        pcd, _ = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
                        if len(pcd.points) > 50:
                            labels = np.array(pcd.cluster_dbscan(eps=0.02, min_points=10))
                            if len(labels) > 0 and np.any(labels >= 0):
                                target_label = np.argmax(np.bincount(labels[labels >= 0]))
                                target_pcd = pcd.select_by_index(np.where(labels == target_label)[0])
                                center = target_pcd.get_center()
                                
                                # 计算物体尺寸
                                width, height, depth = self.calculate_object_dimensions(target_pcd)
                                
                                # 更新标签包含 XYZ 信息（第二行）
                                if width is not None and height is not None and depth is not None:
                                    label_line2 = f"X:{center[0]:.2f} Y:{center[1]:.2f} Z:{center[2]:.2f}"
                                    label_line3 = f"宽:{width:.2f} 高:{height:.2f} 距离:{depth:.2f}"
                                    self.get_logger().info(
                                        f'物体: {obj_name} 中心位置: X={center[0]:.3f}m, Y={center[1]:.3f}m, Z={center[2]:.3f}m, '
                                        f'尺寸: 宽={width:.3f}m, 高={height:.3f}m, 距离={depth:.3f}m'
                                    )
                                else:
                                    label_line2 = f"X:{center[0]:.2f} Y:{center[1]:.2f} Z:{center[2]:.2f}"
                                    label_line3 = None
                                    self.get_logger().info(f'物体: {obj_name} 中心位置: X={center[0]:.3f}m, Y={center[1]:.3f}m, Z={center[2]:.3f}m')

                # 在图像上绘制标签（使用PIL支持中文）
                # 将OpenCV图像转换为PIL图像
                pil_image = PILImage.fromarray(cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB))
                draw = ImageDraw.Draw(pil_image)
                
                # 计算文本位置
                text_y = y1 - 40
                
                # 第一行：物体名称
                draw.text((x1, text_y), label_line1, font=self.font, fill=(0, 165, 255))
                text_y += 15
                
                # 第二行：位置信息
                draw.text((x1, text_y), label_line2, font=self.font, fill=(0, 165, 255))
                text_y += 15
                
                # 第三行：尺寸信息（如果有）
                if label_line3:
                    draw.text((x1, text_y), label_line3, font=self.font, fill=(0, 165, 255))
                
                # 将PIL图像转换回OpenCV格式
                annotated_image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)

            # 发布处理后的图像到 detect 话题
            out_msg = self.bridge.cv2_to_imgmsg(annotated_image, encoding='bgr8')
            out_msg.header = color_msg.header
            self.image_pub.publish(out_msg)

        except Exception as e:
            self.get_logger().error(f'处理失败: {e}')

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

if __name__ == '__main__':
    main()