import os
import time
import threading
import ast
from collections import defaultdict
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 sensor_msgs.msg import Image
from geometry_msgs.msg import Pose, Point, Quaternion
from interfaces.msg import PoseClassAndID, ImgMsg
from interfaces.srv import VisionObjectRecognition


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, depth_img: np.ndarray, intrinsics, hand_eye_mat):
    """计算位态"""
    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=3.0,
    )

    point_cloud = point_cloud.remove_non_finite_points()

    down_pcd = point_cloud.voxel_down_sample(voxel_size=0.020)
    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])
        rmat = hand_eye_mat @ rmat
        rw, rx, ry, rz = tfs.quaternions.mat2quat(rmat[0:3, 0:3])
        x, y, z = rmat[0:3, 3].flatten()

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

    else:
        return 0.0, 0.0, 0.0, None, None, None, None


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(depth_img, mask):
    """
    输入:
        depth_img: H x W
        mask: H x W (0/1 或 bool)
    输出:
        depth_crop, mask_crop
    """
    high, width = depth_img.shape
    ys, xs = np.where(mask > 0)

    if xs.size == 0 or ys.size == 0:
        # 没有有效像素
        return None, None, (None, None)

    x_min, x_max = int(round(xs.min())), int(round(xs.max()))
    y_min, y_max = int(round(ys.min())), int(round(ys.max()))

    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 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.calculate_function = calculate_pose_cpu
        self.fx = self.fy = 0.5
        self.camera_data = {}
        self.cv_bridge = CvBridge()
        self.lock = threading.Lock()
        self.red_range = [
            np.array([0, 120, 70]),
            np.array([10, 255, 255]),
            np.array([170, 120, 70]),
            np.array([180, 255, 255])
        ]

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

    def _init_param(self):
        """init parameter"""
        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('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_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_publisher(self):
        """init publisher"""
        if self.output_boxes or self.output_masks:
            self.pub_detect_image = self.create_publisher(Image, '/image/detect_image', 10)

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

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

    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 = []

                return response
            
        if request.camera_position == 'left' or request.camera_position == 'right':
            hand_eye_mat = self.eye_in_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_dict = self._seg_image(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_dict:
            response.info = "Success get pose"
            response.success = True
            for (class_id, class_name), pose_list in pose_dict.items():
                response.objects.append(
                    PoseClassAndID(
                        class_name = class_name,
                        class_id = class_id,
                        pose_list = pose_list
                    )
                )

            return response
        
        else:
            response.info = "pose dict is empty"
            response.success = False
            response.objects = []

            return response

    def _seg_image(self, rgb_img: np.ndarray, depth_img: np.ndarray, hand_eye_mat):
        """Use segmentation model"""
        pose_dict = defaultdict(list)

        hsv_img = cv2.cvtColor(rgb_img, cv2.COLOR_BGR2HSV)

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

        hsv_img[depth_filter_mask == 0] = 0

        mask1 = cv2.inRange(hsv_img, self.red_range[0], self.red_range[1])
        mask2 = cv2.inRange(hsv_img, self.red_range[2], self.red_range[3])
        mask = mask1 + mask2

        mask = mask // 255

        rgb_img[mask > 0] = rgb_img[mask > 0] * 0.5 + np.array([255, 0, 0]) * 0.5

        orig_shape = depth_img.shape
        mask = cv2.resize(mask.astype(np.uint8), orig_shape[::-1], interpolation=cv2.INTER_NEAREST)
        depth_crop, mask_crop, (x_min, y_min) = crop_mask_bbox(depth_img, mask)

        if depth_crop is None:
            return None, None

        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, depth_crop, intrinsics, hand_eye_mat)

        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_dict[int(98), "red"].append(pose)

        return self.cv_bridge.cv2_to_imgmsg(rgb_img, "bgr8"), pose_dict


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


if __name__ == '__main__':
    main()
