#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
实时检测红色圆点并获取 3D 坐标 (m)。
深度图对齐到彩色图。
"""

import cv2
import numpy as np
import pyrealsense2 as rs
import sys
import time
from pathlib import Path

# 把.so所在目录加入 Python 路径
sys.path.append("/home/lgv/cmvr/cmvr-es/cmake-build-debug/example")

# 导入机器人控制模块
from robot_wrapper import Robot

REPO_ROOT = Path(__file__).resolve().parents[2]
ROBOT_CONFIG = str(REPO_ROOT / "cmvr-es/config/cmvr_es.pb.txt")


def init_realsense():
    """初始化 RealSense 管道并返回 pipeline 和 align 对象"""
    pipeline = rs.pipeline()
    config = rs.config()
    config.enable_stream(rs.stream.depth, 1280, 720, 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, vis_path=None, warmup=30,
                          hough_params=None):
    """
    获取最近红色圆点的 3D 坐标。
    返回 (x, y, z) 或 None（如果没检测到）。
    """
    if hough_params is None:
        hough_params = dict(dp=1.2, minDist=20,
                            param1=50, param2=15,
                            minRadius=5, maxRadius=50)

    # 丢掉前几帧，让相机稳定
    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 并提取红色区域
    hsv = cv2.cvtColor(color_image, cv2.COLOR_BGR2HSV)
    lower_red1, upper_red1 = np.array([0, 100, 100]), np.array([10, 255, 255])
    lower_red2, upper_red2 = np.array([160, 100, 100]), np.array([179, 255, 255])
    mask = cv2.inRange(hsv, lower_red1, upper_red1) | cv2.inRange(hsv, lower_red2, upper_red2)

    # 平滑处理再检测圆
    mask_blur = cv2.GaussianBlur(mask, (9, 9), 2)
    circles = cv2.HoughCircles(mask_blur, cv2.HOUGH_GRADIENT, **hough_params)

    closest_point, min_depth = None, float('inf')

    if circles is not None:
        circles = np.uint16(np.around(circles))
        for u, v, r in circles[0, :]:
            depth = depth_frame.get_distance(int(u), int(v))
            if 0 < depth < min_depth:
                min_depth = depth
                closest_point = rs.rs2_deproject_pixel_to_point(depth_intrin, [int(u), int(v)], depth)
                if vis_path:
                    cv2.circle(color_image, (u, v), r, (0, 255, 0), 2)
                    cv2.circle(color_image, (u, v), 6, (0, 0, 255), -1)

    if vis_path and closest_point is not None:
        cv2.imwrite(vis_path, color_image)

    return closest_point


def main(serial="243122075614", avg_window=3, loop=True):
    pipeline, align = init_realsense()
    # robot = Robot(ROBOT_CONFIG, "hc01")

    try:
        while True:
            point_3d = get_closest_red_point(pipeline, align, vis_path="red_point.png")
            if point_3d:
                print("Closest red 3D point:", np.round(point_3d, 4))
            else:
                print("未检测到红点")

            if not loop:
                break

            key = cv2.waitKey(1)
            if key == 27:  # ESC 键退出
                break

    except Exception as e:
        print("Error:", e)

    finally:
        pipeline.stop()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()