#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
import numpy as np
import cv2
import tf2_ros
import time
from rclpy.time import Time

"""
==============================
手眼标定数据采集节点（完整版本）
Eye-to-Hand
==============================
"""
data = np.load("/home/shen/openarm/openarm_test/src/openarm_moveit_control/openarm_moveit_control/stereo_params_1920x1080.npz")
mtxL = data["mtxL"]
distL = data["distL"]
mtxR = data["mtxR"]
distR = data["distR"]
R = data["R"]
T = data["T"]
image_size = tuple(data["image_size"])

R1, R2, P1, P2, Q0, roi1, roi2 = cv2.stereoRectify(
    mtxL, distL,
    mtxR, distR,
    image_size,
    R, T,
    flags=cv2.CALIB_ZERO_DISPARITY,
    alpha=0
)

mapL1, mapL2 = cv2.initUndistortRectifyMap(
    mtxL, distL, R1, P1, image_size, cv2.CV_32FC1
)
mapR1, mapR2 = cv2.initUndistortRectifyMap(
    mtxR, distR, R2, P2, image_size, cv2.CV_32FC1
)
print("Q:", Q0)

class HandEyeCollector(Node):

    def __init__(self):
        super().__init__('handeye_collector')

        # ========= TF =========
        self.tf_buffer = tf2_ros.Buffer()
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer, self)

        # ========= 相机 =========
        self.cap = cv2.VideoCapture(15)  # ← 改成你的相机 ID
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 2560)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
        self.cap.set(cv2.CAP_PROP_FPS, 30)
        self.cap.set(cv2.CAP_PROP_EXPOSURE, -8)
        

        # ========= ArUco =========
        self.aruco_dict = cv2.aruco.getPredefinedDictionary(
            cv2.aruco.DICT_4X4_50
        )
        self.aruco_params = cv2.aruco.DetectorParameters()

        # ========= 数据 =========
        self.cam_points = []
        self.base_points = []

        # ========= 双目标定参数（必须填） =========
        # 方法 A：Q 矩阵（推荐）
        self.Q = Q0   # ← TODO: 填你的 4x4 Q 矩阵

        # 方法 B：手动参数（二选一）
        self.fx = None
        self.fy = None
        self.cx = None
        self.cy = None
        self.baseline = None  # 单位：米
        self.get_logger().info(
            "\nHand-eye collector ready\n"
            "ENTER: save sample\n"
            "ESC: save & quit\n"
        )
        self.create_timer(0.5, self.spin)
        # self.spin()

    # ===============================
    # 工具函数
    # ===============================
    def check(self):
        try:
            tf = self.tf_buffer.lookup_transform(
                'openarm_body_link0',
                'openarm_left_link7',
                Time()
            )
            print("Transform:", tf)
            self.get_logger().info("TF OK")
        except:
            self.get_logger().warn("TF not available yet")
            return None

    def get_tcp_position(self):
        try:
            # if not self.tf_buffer.can_transform(
            #     'openarm_body_link0',
            #     'openarm_left_link7',
            #     Time()
            # ):
            #     self.get_logger().warn("TF not available yet")
            #     return None

            t = self.tf_buffer.lookup_transform(
                'openarm_body_link0',
                'openarm_left_link7',
                Time()
            )
            print("Transform:", t)
            return np.array([
                t.transform.translation.x,
                t.transform.translation.y,
                t.transform.translation.z
            ])
        except Exception as e:
            self.get_logger().warn(f"TF lookup failed: {str(e)}")
            return None


    def split_stereo(self, frame):
        """2560x720 → 左右 1280x720"""
        h, w, _ = frame.shape
        left = frame[:, :w//2]
        right = frame[:, w//2:]
        return left, right

    def compute_disparity(self, left, right):
        gray_l = cv2.cvtColor(left, cv2.COLOR_BGR2GRAY)
        gray_r = cv2.cvtColor(right, cv2.COLOR_BGR2GRAY)

        stereo = cv2.StereoSGBM_create(
            minDisparity=0,
            numDisparities=128,
            blockSize=7
        )
        disp = stereo.compute(gray_l, gray_r).astype(np.float32) / 16.0
        return disp

    def pixel_to_3d(self, u, v, disparity):
        """
        像素 + 视差 → 3D（camera_link）
        """
        if self.Q is not None:
            point_4d = np.array([u, v, disparity, 1.0])
            X = self.Q @ point_4d
            X /= X[3]
            return X[:3]

    # ===============================
    # 主循环
    # ===============================

    def spin(self):
        
        # frame_count = 0
        while rclpy.ok():
            ret, frame = self.cap.read()
            if not ret:
                continue

            left, right = self.split_stereo(frame)
            rectL = cv2.remap(left,  mapL1, mapL2, cv2.INTER_LINEAR)
            rectR = cv2.remap(right, mapR1, mapR2, cv2.INTER_LINEAR)
            # frame_count += 1
            # if frame_count % 3 == 0:
            small_l = cv2.resize(rectL, None, fx=0.5, fy=0.5)
            small_r = cv2.resize(rectR, None, fx=0.5, fy=0.5)
            disp = self.compute_disparity(small_l, small_r)
            gray = cv2.cvtColor(rectL, cv2.COLOR_BGR2GRAY)
            corners, ids, _ = cv2.aruco.detectMarkers(
                gray, self.aruco_dict, parameters=self.aruco_params
            )

            if ids is not None:
                cv2.aruco.drawDetectedMarkers(gray, corners, ids)
                c = corners[0][0]
                center = c.mean(axis=0)
                u, v = int(center[0]), int(center[1])
                d = disp[v, u]

                cv2.circle(rectL, (u, v), 5, (0, 0, 255), -1)
                cv2.putText(
                    rectL, f"disp={d:.2f}", (u+5, v-5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 1
                )

            cv2.imshow("left", rectL)
            # self.check()
            key = cv2.waitKey(1)

            # ENTER：保存
            if key == 13 and ids is not None:
                cam_p = self.pixel_to_3d(u, v, d)
                print("cam_p:", cam_p)
                base_p = self.get_tcp_position()
                if cam_p is None or base_p is None:
                    self.get_logger().warn("Invalid sample")
                    continue

                self.cam_points.append(cam_p)
                self.base_points.append(base_p)

                self.get_logger().info(
                    f"Saved #{len(self.cam_points)} "
                    f"cam={cam_p.round(3)} "
                    f"base={base_p.round(3)}"
                )
                time.sleep(0.3)

            # ESC：退出
            elif key == 27:
                break

        self.save_and_exit()

    def save_and_exit(self):
        np.savez(
            "handeye_data.npz",
            cam=np.array(self.cam_points),
            base=np.array(self.base_points)
        )
        self.get_logger().info(
            f"Saved {len(self.cam_points)} samples to handeye_data.npz"
        )
        self.cap.release()
        cv2.destroyAllWindows()


# ===============================
# main
# ===============================

def main():
    rclpy.init()
    node = HandEyeCollector()
    rclpy.spin(node)
    rclpy.shutdown()

if __name__ == "__main__":
    main()
