#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
从 RealSense 相机中获取一帧对齐后的彩色 / 深度图，
检测棋盘格左上角角点，并反投影到相机坐标系 (m)。
"""

import cv2
import time
import numpy as np
import pyrealsense2 as rs


# ========================= 1. 管线启动 / 关闭 ============================== #
def start_pipeline(serial: str, depth_size=(640, 480), color_size=(640, 480), fps: int = 30):
    pipeline = rs.pipeline()
    cfg = rs.config()
    cfg.enable_device(serial)

    # 正确的参数顺序:  width, height, format, fps
    cfg.enable_stream(rs.stream.depth, depth_size[0], depth_size[1], rs.format.z16, fps)
    cfg.enable_stream(rs.stream.color, color_size[0], color_size[1], rs.format.bgr8, fps)

    pipeline_profile = pipeline.start(cfg)
    depth_intr = pipeline_profile.get_stream(rs.stream.depth).as_video_stream_profile().get_intrinsics()
    align = rs.align(rs.stream.depth)
    return pipeline, align, depth_intr

def stop_pipeline(pipeline: rs.pipeline):
    pipeline.stop()


# ========================= 2. 获取对齐后的单帧 ============================ #
def get_aligned_frames(pipeline: rs.pipeline, align: rs.align, warmup: int = 10, timeout_ms: int = 5000):
    for _ in range(warmup):
        pipeline.wait_for_frames()

    # 等待一帧
    for _ in range(3):                      # 最多重试 3 次
        try:
            frames = pipeline.wait_for_frames(timeout_ms=timeout_ms)
            aligned = align.process(frames)
            return aligned.get_color_frame(), aligned.get_depth_frame()
        except RuntimeError:
            print("⚠️ 等待帧超时，重试中...")
            time.sleep(1)
    raise RuntimeError("连续超时，无法获取有效帧")


# ========================= 3. 计算目标 3-D 坐标 =========================== #
def chessboard_lefttop_position(color_frame, depth_frame, depth_intr, pattern_size=(8, 11), avg_window=3, vis_path=None):        # ← 额外参数：保存路径
    """
    返回左上角角点相机坐标 (X,Y,Z) [m]。
    若 vis_path 不为 None，则在彩色图上标出中心和角点并保存到该路径。
    """
    color_img = np.asanyarray(color_frame.get_data()).copy()
    gray = cv2.cvtColor(color_img, cv2.COLOR_BGR2GRAY)

    ok, corners = cv2.findChessboardCorners(
        gray, pattern_size,
        flags=cv2.CALIB_CB_ADAPTIVE_THRESH | cv2.CALIB_CB_NORMALIZE_IMAGE
    )
    if not ok:
        raise RuntimeError("未检测到棋盘格")

    cv2.cornerSubPix(
        gray, corners, winSize=(5, 5), zeroZone=(-1, -1),
        criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_MAX_ITER, 30, 1e-3)
    )

    # --------- 取第 0 个角点像素坐标 ---------
    u, v = corners[0].ravel()

    # --------- 可视化：画中心十字 & 左上角红点 ---------
    if vis_path is not None:
        cx, cy = depth_intr.ppx, depth_intr.ppy                # 主点
        cv2.drawMarker(color_img, (int(cx), int(cy)), (0, 255, 0),
                       markerType=cv2.MARKER_CROSS, markerSize=20, thickness=2)
        cv2.circle(color_img, (int(u), int(v)), 6, (0, 0, 255), -1)
        cv2.imwrite(vis_path, color_img)                       # ← 保存文件
        print(f"✓ 已保存标记图到 {vis_path}")

    # … 以下深度中值滤波 + 反投影保持不变 … ----------------------------
    width, height = depth_frame.get_width(), depth_frame.get_height()
    depths = [depth_frame.get_distance(int(round(u + du)), int(round(v + dv)))
              for du in range(-avg_window, avg_window + 1)
              for dv in range(-avg_window, avg_window + 1)
              if 0 <= int(round(u + du)) < width and
              0 <= int(round(v + dv)) < height and
              depth_frame.get_distance(int(round(u + du)), int(round(v + dv))) > 0]

    if not depths:
        raise RuntimeError("角点深度无效")
    depth = float(np.median(depths))
    X, Y, Z = rs.rs2_deproject_pixel_to_point(depth_intr, [u, v], depth)
    return np.array([X, Y, Z], dtype=np.float32)


# ========================= 4. 主程序入口 ================================= #
def main(serial="243122075614", pattern_size=(8, 11), avg_window=3):
    pipeline, align, depth_intr = start_pipeline(serial)
    try:
        color_f, depth_f = get_aligned_frames(pipeline, align)
        pos = chessboard_lefttop_position(
            color_f, depth_f, depth_intr,
            pattern_size=pattern_size, avg_window=avg_window,
            vis_path="chessboard_marked.png"
        )
        print("左上角 3-D 坐标 (m):", pos)
    finally:
        stop_pipeline(pipeline)


if __name__ == "__main__":
    main()