import pyrealsense2 as rs
import numpy as np
import cv2

def init_realsense():
    """
    初始化 RealSense 管道并返回 pipeline 和 align 对象
    """
    pipeline = rs.pipeline()
    config = rs.config()
    config.enable_stream(rs.stream.depth, 848, 480, rs.format.z16, 30)
    config.enable_stream(rs.stream.color, 1280, 720, rs.format.bgr8, 30)
    profile = pipeline.start(config)
    align = rs.align(rs.stream.color)
    return pipeline, align

def get_closest_red_point(pipeline, align, warmup=100):
    """
    获取最近红色圆点的 3D 坐标。
    返回 (x, y, z) 或 None（如果没检测到）。
    """
    # 丢掉前几帧，让相机稳定
    for _ in range(warmup):
        pipeline.wait_for_frames()

    frames = pipeline.wait_for_frames()
    aligned_frames = align.process(frames)

    depth_frame = aligned_frames.get_depth_frame()
    color_frame = aligned_frames.get_color_frame()
    if not depth_frame or not color_frame:
        return None

    depth_image = np.asanyarray(depth_frame.get_data())
    color_image = np.asanyarray(color_frame.get_data())
    depth_intrin = depth_frame.profile.as_video_stream_profile().intrinsics

    hsv = cv2.cvtColor(color_image, cv2.COLOR_BGR2HSV)
    lower_red1 = np.array([0, 100, 100])
    upper_red1 = np.array([10, 255, 255])
    lower_red2 = np.array([160, 100, 100])
    upper_red2 = np.array([179, 255, 255])
    mask1 = cv2.inRange(hsv, lower_red1, upper_red1)
    mask2 = cv2.inRange(hsv, lower_red2, upper_red2)
    mask = cv2.bitwise_or(mask1, mask2)

    mask_blur = cv2.GaussianBlur(mask, (9, 9), 2)
    circles = cv2.HoughCircles(mask_blur, cv2.HOUGH_GRADIENT, dp=1.2, minDist=20,
                               param1=50, param2=15, minRadius=5, maxRadius=50)

    closest_point = None
    min_depth = float('inf')

    if circles is not None:
        circles = np.uint16(np.around(circles))
        for i in circles[0, :]:
            u, v, r = i
            depth = depth_frame.get_distance(u, v)
            if 0 < depth < min_depth:
                min_depth = depth
                closest_point = rs.rs2_deproject_pixel_to_point(depth_intrin, [u, v], depth)

    return closest_point

def show_frames(pipeline, align):
    """
    显示 RGB 和 Depth 图像，同时标记最近红点（实时窗口）。
    按 'q' 退出。
    """
    while True:
        frames = pipeline.wait_for_frames()
        aligned_frames = align.process(frames)

        depth_frame = aligned_frames.get_depth_frame()
        color_frame = aligned_frames.get_color_frame()
        if not depth_frame or not color_frame:
            continue

        depth_image = np.asanyarray(depth_frame.get_data())
        color_image = np.asanyarray(color_frame.get_data())
        depth_intrin = depth_frame.profile.as_video_stream_profile().intrinsics

        hsv = cv2.cvtColor(color_image, cv2.COLOR_BGR2HSV)
        lower_red1 = np.array([0, 100, 100])
        upper_red1 = np.array([10, 255, 255])
        lower_red2 = np.array([160, 100, 100])
        upper_red2 = np.array([179, 255, 255])
        mask1 = cv2.inRange(hsv, lower_red1, upper_red1)
        mask2 = cv2.inRange(hsv, lower_red2, upper_red2)
        mask = cv2.bitwise_or(mask1, mask2)

        mask_blur = cv2.GaussianBlur(mask, (9, 9), 2)
        circles = cv2.HoughCircles(mask_blur, cv2.HOUGH_GRADIENT, dp=1.2, minDist=20,
                                   param1=50, param2=15, minRadius=5, maxRadius=50)

        closest_point = None
        closest_pixel = None
        closest_radius = 0
        min_depth = float('inf')

        if circles is not None:
            circles = np.uint16(np.around(circles))
            for i in circles[0, :]:
                u, v, r = i
                depth = depth_frame.get_distance(u, v)
                if 0 < depth < min_depth:
                    min_depth = depth
                    closest_point = rs.rs2_deproject_pixel_to_point(depth_intrin, [u, v], depth)
                    closest_pixel = (u, v)
                    closest_radius = r

        if closest_point is not None:
            cv2.circle(color_image, closest_pixel, closest_radius, (0, 255, 0), 2)
            cv2.circle(color_image, closest_pixel, 2, (255, 0, 0), 3)

        cv2.imshow("RGB Image with Closest Red Circle", color_image)
        depth_colormap = cv2.convertScaleAbs(depth_image, alpha=0.03)
        depth_colormap = cv2.applyColorMap(depth_colormap, cv2.COLORMAP_JET)
        cv2.imshow("Depth Image", depth_colormap)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

if __name__ == "__main__":
    pipeline, align = init_realsense()
    try:
        # 调用获取最近红点
        point_3d = get_closest_red_point(pipeline, align)
        print("Closest red 3D point:", point_3d)

        # 调用显示函数（实时窗口显示）
        show_frames(pipeline, align)
    finally:
        pipeline.stop()
        cv2.destroyAllWindows()