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 draw_pointcloud(pcd, T):

    R = T[0:3, 0:3]
    point = T[0:3, 3:4].flatten()
    x, y, z = R[:, 0], R[:, 1], R[:, 2]

    points = [
        [0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1],
        point, point + x, point + y, point + z

    ]  # 画点：原点、第一主成分、第二主成分
    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]
    ]

    line_set = o3d.geometry.LineSet(points=o3d.utility.Vector3dVector(points), lines=o3d.utility.Vector2iVector(lines))
    line_set.colors = o3d.utility.Vector3dVector(colors)

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


def preprocess_point_cloud(pcd, voxel_size):

    pcd_down = pcd.voxel_down_sample(voxel_size)
    radius_normal = voxel_size * 2

    pcd_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=radius_normal, max_nn=30))

    radius_feature = voxel_size * 5

    pcd_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
        pcd_down,
        o3d.geometry.KDTreeSearchParamHybrid(radius=radius_feature, max_nn=100)
    )

    return pcd_down, pcd_fpfh

def prepare_dataset(source, target, voxel_size):

    trans_init = np.identity(4)
    source.transform(trans_init)

    source_down, source_fpfh = preprocess_point_cloud(source, voxel_size)
    target_down, target_fpfh = preprocess_point_cloud(target, voxel_size)

    return source_down, target_down, source_fpfh, target_fpfh


def execute_global_registration(source_down, target_down, source_fpfh, target_fpfh, voxel_size):

    distance_threshold = voxel_size * 1.5
    result = o3d.pipelines.registration.registration_ransac_based_on_feature_matching(
        source_down,
        target_down,
        source_fpfh,
        target_fpfh,
        True,
        distance_threshold,
        o3d.pipelines.registration.TransformationEstimationPointToPoint(False),
        3,
        [
            o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9),
            o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(distance_threshold)
        ],
        o3d.pipelines.registration.RANSACConvergenceCriteria(100000, 0.999))

    return result.transformation


def object_icp(
        target: o3d.geometry.PointCloud,
        source: o3d.geometry.PointCloud | str,
        ransac_voxel_size: float = 0.005,
        icp_voxel_radius: list[float] | None = None,
        icp_max_iter: list[int] | None = None,
):

    if icp_voxel_radius is None:
        icp_voxel_radius = [0.004, 0.002, 0.001]

    if icp_max_iter is None:
        icp_max_iter = [50, 30, 14]

    if isinstance(source, str):
        source = o3d.io.read_point_cloud(source)
    elif isinstance(source, o3d.geometry.PointCloud):
        pass
    else:
        raise TypeError(f"Unsupported Type {type(source)}")

    voxel_size = 0.005  # means 5mm for this dataset
    source_down, target_down, source_fpfh, target_fpfh = prepare_dataset(source, target, voxel_size)

    T = execute_global_registration(
        source_down,
        target_down,
        source_fpfh,
        target_fpfh,
        ransac_voxel_size
    )

    for scale in range(3):
        iter = icp_max_iter[scale]
        radius = icp_voxel_radius[scale]
        # print([iter, radius, scale])

        source_down = source.voxel_down_sample(radius)
        target_down = target.voxel_down_sample(radius)

        source_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=radius * 2, max_nn=30))
        target_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=radius * 2, max_nn=30))

        result_icp = o3d.pipelines.registration.registration_icp(
            source_down,
            target_down,
            radius * 5,
            T,
            o3d.pipelines.registration.TransformationEstimationPointToPoint(),
            o3d.pipelines.registration.ICPConvergenceCriteria(
                relative_fitness=1e-6,
                relative_rmse=1e-6,
                max_iteration=iter
            )
        )

        T = result_icp.transformation

    draw_pointcloud([source.transform(T), target], T)
    
    return T


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, model_pcd):
    """计算位态"""
    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=2.0,
    )

    point_cloud = point_cloud.remove_non_finite_points()

    clean_pcd, _ = point_cloud.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
    
    rmat = object_icp(model_pcd, clean_pcd)

    x, y, z = rmat[0:3, 3:4].flatten()

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

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

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


def calculate_grav_width(mask, K, depth):
    """计算重心宽度"""
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if contours:
        box = cv2.boxPoints(cv2.minAreaRect(contours[0]))
        if np.linalg.norm(box[1] - box[0]) < np.linalg.norm(box[1] - box[2]):
            point_diff = box[1] - box[0]
        else:
            point_diff = box[1] - box[2]

        grab_width = depth * np.sqrt(point_diff[0]**2 / K[0]**2 + point_diff[1]**2 / K[4]**2)
        return grab_width
    else:
        return 0.0


class DetectNode(Node):
    def __init__(self, name, mode):
        super().__init__(name)
        self.mode = mode
        self.device = None
        self.checkpoint_path = None
        self.checkpoint_name = None
        self.output_boxes = None
        self.output_masks = None
        self.function = None
        self.calculate_function = None

        self.K = None

        self.fx = self.fy = 0.5
        self.camera_data = {}

        self.cv_bridge = CvBridge()
        self.lock = threading.Lock()
        
        '''init'''
        self._init_param()

        self.source_model = o3d.io.read_point_cloud(self.source_model_path)

        if mode == 'detect':
            self._init_model()
        else:
            self.function = None
            self.get_logger().error('Error: Mode Unknown')

        if self.device == 'cpu':
            self.calculate_function = calculate_pose_cpu
        elif self.device == 'gpu':
            raise NotImplementedError('Function: calculate_pose_gpu not implemented')
        else:
            raise ValueError(f"device must be cpu or gpu, now {self.device}")

        self._init_publisher()
        self._init_subscriber()
        self._init_service()

    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.source_model_path = os.path.join(share_dir, 'pointclouds/bottle_model.pcd')

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

        self.declare_parameter('classes', 'None')
        self.classes = ast.literal_eval(self.get_parameter('classes').value)

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

        """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
            # self.get_logger().info('get_pose')
            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
                    )
                )
        
        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):
        """Use segmentation model"""
        pose_dict = defaultdict(list)

        '''Get Predict Results'''
        # time1 = time.time()
        results = self.model(rgb_img, retina_masks=True, conf=self.set_confidence, classes=[39])
        # 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.data.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)
            depth_crop, mask_crop, (x_min, y_min) = crop_mask_bbox(depth_img, mask, box)
            
            if depth_crop is None:
                continue

            if mask.shape[0] >= (orig_shape[0] * 0.5) and mask.shape[1] >= (orig_shape[1] * 0.5):
                mask_crop = cv2.resize(mask_crop, None, fx=self.fx, fy=self.fy, interpolation=cv2.INTER_NEAREST)
                depth_crop = cv2.resize(depth_crop, None, fx=self.fx, fy=self.fy, interpolation=cv2.INTER_NEAREST)
                intrinsics = o3d.camera.PinholeCameraIntrinsic(
                int(self.camera_size[0] * self.fx),
                int(self.camera_size[1] * self.fy),
                self.K[0] * self.fx,
                self.K[4] * self.fy,
                (self.K[2] - x_min) * self.fx,
                (self.K[5] - y_min) * self.fy
            )
            else:
                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, self.source_model)

            grab_width = calculate_grav_width(mask, self.K, z)
            z = z + grab_width * 0.45

            if (x, y, z) == (None, None, None):
                self.get_logger().error("have wrong pose")
                continue

            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(class_ids[i]), labels[class_ids[i]]].append(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_dict
        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_dict
        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_dict
        else:
            return None, pose_dict


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


if __name__ == '__main__':
    main()
