import logging

import cv2
import numpy as np
import open3d as o3d


def draw_box(rgb_img: np.ndarray, segmentation_result, put_text: bool = True):
    """
    绘制目标检测边界框
    """
    boxes = segmentation_result.boxes.xywh.cpu().numpy()
    confidences = segmentation_result.boxes.conf.cpu().numpy()
    class_ids = segmentation_result.boxes.cls.cpu().numpy()
    labels = segmentation_result.names

    for i, box in enumerate(boxes):
        x_center, y_center, width, height = box[:4]

        p1 = [int((x_center - width / 2)), int((y_center - height / 2))]
        p2 = [int((x_center + width / 2)), int((y_center + height / 2))]
        cv2.rectangle(rgb_img, p1, p2, (255, 255, 0), 2)

        if put_text:
            cv2.putText(rgb_img, f'{labels[class_ids[i]]}: {confidences[i]*100:.2f}', (p1[0], p1[1] - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 0), 2)


def draw_mask(rgb_img: np.ndarray, segmentation_result):
    """
    绘制实例分割mask
    """
    masks = segmentation_result.masks.data.cpu().numpy()
    orig_shape = segmentation_result.masks.orig_shape

    for i, mask in enumerate(masks):
        mask = cv2.resize(mask.astype(np.uint8), orig_shape[::-1], interpolation=cv2.INTER_NEAREST)
        rgb_img[mask > 0] = rgb_img[mask > 0] * 0.5 + np.array([0, 0, 255]) * 0.5


def draw_pointcloud(pcd, axis: bool = True):
    """绘制点云"""
    if not pcd:
        logging.warning("pcd is empty")
    if axis:
        axis_point = [
        [0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]
        ]  # 画点：原点、第一主成分、第二主成分
        axis = [
            [0, 1], [0, 2], [0, 3]
        ]  # 画出三点之间两两连线
        axis_colors = [
            [1, 0, 0], [0, 1, 0], [0, 0, 1]
        ]
        # 构造open3d中的LineSet对象，用于主成分显示
        axis_set = o3d.geometry.LineSet(points=o3d.utility.Vector3dVector(axis_point), lines=o3d.utility.Vector2iVector(axis))
        axis_set.colors = o3d.utility.Vector3dVector(axis_colors)

        pcd.append(axis_set)
    o3d.visualization.draw_geometries(pcd)
