import rclpy
from rclpy.node import Node
from rclpy.action import ActionClient
import numpy as np
import time
import cv2

from moveit_msgs.action import MoveGroup
from moveit_msgs.msg import Constraints, PositionConstraint
from geometry_msgs.msg import PoseStamped
from shape_msgs.msg import SolidPrimitive

data = np.load("/home/shen/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"])

handey_cal = np.load("/home/shen/openarm_test/handeye_data.npz")

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 HandFollowMoveIt(Node):
    def __init__(self):
        super().__init__('hand_follow_moveit')

        self.client = ActionClient(self, MoveGroup, '/move_action')

        self.last_sent = time.time()
        self.send_period = 0.3   # 约 3 Hz

        self.p_prev = None

        # ========= 相机 =========
        self.cap = cv2.VideoCapture(0)  # ← 改成你的相机 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()
        # ========= 双目标定参数（必须填） =========
        # 方法 A：Q 矩阵（推荐）
        self.Q = Q0   # ← TODO: 填你的 4x4 Q 矩阵
        self.create_timer(0.5, self.spin)
    # ==============================
    # 手眼标定求解
    # ==============================
    def solve_handeye(self, cam, base):
        cA = cam.mean(axis=0)
        cB = base.mean(axis=0)
        AA = cam - cA
        BB = base - cB

        U, _, Vt = np.linalg.svd(AA.T @ BB)
        R = Vt.T @ U.T
        if np.linalg.det(R) < 0:
            Vt[2] *= -1
            R = Vt.T @ U.T

        t = cB - R @ cA
        return R, t
    
    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 update_hand(self, p):
        """ p: np.array([x,y,z]) in base_link """

        # 1️⃣ 限幅（按你机械臂改）
        p[0] = np.clip(p[0], 0.25, 0.55)
        p[1] = np.clip(p[1], -0.30, 0.30)
        p[2] = np.clip(p[2], 0.25, 0.30)

        # 2️⃣ 低通滤波
        if self.p_prev is None:
            self.p_prev = p
            return

        alpha = 0.2
        p = alpha * p + (1 - alpha) * self.p_prev

        # 3️⃣ 死区
        if np.linalg.norm(p - self.p_prev) < 0.015:
            return

        # 4️⃣ 控制频率
        if time.time() - self.last_sent < self.send_period:
            return

        self.send_goal(p)
        self.p_prev = p
        self.last_sent = time.time()

    def send_goal(self, p):
        pose = PoseStamped()
        pose.header.frame_id = "openarm_body_link0"
        pose.pose.position.x = float(p[0])
        pose.pose.position.y = float(p[1])
        pose.pose.position.z = float(p[2])
        pose.pose.orientation.w = 1.0

        goal = MoveGroup.Goal()
        req = goal.request
        req.group_name = "left_arm"
        req.allowed_planning_time = 2.0
        req.max_velocity_scaling_factor = 0.2

        pc = PositionConstraint()
        pc.header = pose.header
        pc.link_name = "openarm_left_link7"

        box = SolidPrimitive()
        box.type = SolidPrimitive.BOX
        box.dimensions = [0.02, 0.02, 0.02]

        pc.constraint_region.primitives.append(box)
        pc.constraint_region.primitive_poses.append(pose.pose)

        constraints = Constraints()
        constraints.position_constraints.append(pc)
        req.goal_constraints.append(constraints)

        self.client.wait_for_server()
        self.client.send_goal_async(goal)

    def get_hand_3d(self):
        return None
    
    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(rectL, rectR)
                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
                    )
                    self.solve_handeye(handey_cal["cam"], handey_cal["base"])

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


def main():
    rclpy.init()
    node = HandFollowMoveIt()
    rclpy.spin(node)
    rclpy.shutdown()
    # while True:
    #     p_hand_base = get_hand_3d()   # 你已有
    #     node.update_hand(p_hand_base)

if __name__ == "__main__":
    main()