import os
import time
from ament_index_python.packages import get_package_share_directory

import cv2
import open3d as o3d
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 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


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 pca(data, sort=True):
    """主成分分析 """
    center = np.mean(data, axis=0)
    centered_points = data - center    # 去中心化

    try:
        cov_matrix = np.cov(centered_points.T)  # 转置
        eigenvalues, eigenvectors = np.linalg.eigh(cov_matrix)

    except np.linalg.LinAlgError:
        return None, None

    if sort:
        sort = eigenvalues.argsort()[::-1]  # 降序排列
        eigenvalues = eigenvalues[sort]  # 索引
        eigenvectors = eigenvectors[:, sort]

    return eigenvalues, eigenvectors

def calculate_pose_cpu(mask, rgb_img, depth_img: np.ndarray, intrinsics):
    """计算位态"""
    depth_img_mask = np.zeros_like(depth_img)
    depth_img_mask[mask > 0] = depth_img[mask > 0]

    depth_o3d = o3d.geometry.Image(depth_img_mask.astype(np.uint16))

    point_cloud = o3d.geometry.PointCloud.create_from_depth_image(
        depth=depth_o3d,
        intrinsic=intrinsics,
        depth_scale=1000.0,
        depth_trunc=8.0,
    )

    point_cloud = point_cloud.remove_non_finite_points()

    down_pcd = point_cloud.voxel_down_sample(voxel_size=0.022)
    clean_pcd, _ = down_pcd.remove_radius_outlier(nb_points=10, radius=0.1)
    clean_pcd, _ = clean_pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=3.0)

    if len(clean_pcd.points) == 0:
        return 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
    center = clean_pcd.get_center()
    x, y, z = center

    w, v = pca(np.asarray(clean_pcd.points))

    if w is not None and v is not None:
        vx, vy, vz = v[:,0], v[:,1], v[:,2]

        if vx[0] < 0:
            vx = -vx
        if vy[1] < 0:
            vy = -vy
        if not np.allclose(np.cross(vx, vy), vz):
            vz = -vz

        R = np.column_stack((vx, vy, vz))
        rmat = tfs.affines.compose(np.squeeze(np.asarray((x, y, z))), R, [1, 1, 1])

        rw, rx, ry, rz = tfs.quaternions.mat2quat(rmat[0:3, 0:3])
        x, y, z = rmat[0:3, 3].flatten()

        # point = [
        #     [x, y, z], [x, y, z] + vx, [x, y, z] + vy, [x, y, z] + vz,
        #     [0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]
        #     ]  # 画点：原点、第一主成分、第二主成分
        # lines = [
        #     [0, 1], [0, 2], [0, 3], [4, 5], [4, 6], [4, 7]
        # ]  # 画出三点之间两两连线
        # colors = [
        #     [1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 0, 0], [0, 1, 0], [0, 0, 1]
        # ]
        # # 构造open3d中的LineSet对象，用于主成分显示
        # line_set = o3d.geometry.LineSet(points=o3d.utility.Vector3dVector(point), lines=o3d.utility.Vector2iVector(lines))
        # line_set.colors = o3d.utility.Vector3dVector(colors)
        # o3d.visualization.draw_geometries([point_cloud, line_set])
        # o3d.visualization.draw_geometries([clean_pcd, line_set])

        return x, y, z, rw, rx, ry, rz


def draw_box(set_confidence, rgb_img, result):
    """绘制目标检测边界框"""
    boxes = result.boxes.xywh.cpu().numpy()
    confidences = result.boxes.conf.cpu().numpy()
    class_ids = result.boxes.cls.cpu().numpy()
    labels = result.names

    for i, box in enumerate(boxes):
        if confidences[i] >= set_confidence:
            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)

            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(set_confidence, rgb_img, result):
    """绘制实例分割mask"""
    masks = result.masks.data.cpu().numpy()
    orig_shape = result.masks.orig_shape
    confidences = result.boxes.conf.cpu().numpy()

    for i, mask in enumerate(masks):
        if confidences[i] >= set_confidence:
            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
        else:
            continue


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.eye_in_hand_mat = None
        self.eye_to_hand_mat = None

        self.calculate_function = calculate_pose_cpu
        self.fx = self.fy = 1.0
        self.cv_bridge = CvBridge()

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

    def _init_param(self):
        """init parameter"""
        self.declare_parameter('checkpoint_name', 'hivecorebox-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', True)
        self.output_masks = self.get_parameter('output_masks').value

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

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

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

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

        self.declare_parameter('eye_in_hand', [1.0, 0.0, 0.0, 0.0,
                                              0.0, 1.0, 0.0, 0.0,
                                              0.0, 0.0, 1.0, 0.0,
                                              0.0, 0.0, 0.0, 1.0])
        self.eye_in_hand_mat = np.array(self.get_parameter('eye_in_hand').value).reshape(4, 4)

        self.declare_parameter('eye_to_hand', [1.0, 0.0, 0.0, 0.0,
                                              0.0, 1.0, 0.0, 0.0,
                                              0.0, 0.0, 1.0, 0.0,
                                              0.0, 0.0, 0.0, 1.0])
        self.eye_to_hand_mat = np.array(self.get_parameter('eye_to_hand').value).reshape(4, 4)


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

        '''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_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"]
                    )
                )
            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_list = []

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

        rgb_img[depth_filter_mask == 0] = 0

        '''Get Predict Results'''
        time1 = time.time()
        results = self.model(rgb_img, retina_masks=True, 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

            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, rgb_crop, depth_crop, intrinsics)
            print(x, y, z)

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

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

        '''mask_img and box_img is or not output'''
        if self.output_boxes and not self.output_masks:
            draw_box(self.set_confidence, 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(self.set_confidence, rgb_img, result)
            draw_mask(self.set_confidence, 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(self.set_confidence, rgb_img, result)
            return self.cv_bridge.cv2_to_imgmsg(rgb_img, "bgr8"), pose_list
        else:
            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()
