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

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

import tf2_ros
import tf2_geometry_msgs



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.plan_period = 1.0       # 秒（规划频率）
        self.min_delta = 0.01        # m，手移动超过才重新规划
        self.z_min = 0.15            # 防止撞桌
        self.z_max = 0.60
        self.latest_hand = None
        self.last_target = None

        self.base_frame = 'openarm_body_link0'

        # ====== TF ======
        self.tf_buffer = tf2_ros.Buffer()
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer, self)
        self.send_goal([0.0, 0.15, 0.3])
        time.sleep(3)

        # ====== 订阅 ======
        self.sub = self.create_subscription(
            PointStamped,
            '/hand_3d',
            self.hand_cb,
            10
        )
        self.timer = self.create_timer(self.plan_period, self.plan_timer)

        self.get_logger().info('Hand → MoveIt Planning node started')
    # ---------- 回调 ----------
    def hand_cb(self, msg):
        try:
            hand = self.tf_buffer.transform(
                msg,
                self.base_frame,
                timeout=rclpy.duration.Duration(seconds=0.1)
            )
        except Exception:
            return
        print(f"手坐标（{self.base_frame}):{hand.point.x:.3f}, {hand.point.y:.3f}, {hand.point.z:.3f}")
        self.latest_hand = np.array([
            hand.point.x,
            hand.point.y,
            hand.point.z
        ])

    def update_hand(self, p):
        """ p: np.array([x,y,z]) in base_link """

        # 1️⃣ 限幅（按你机械臂改）
        p[0] = np.clip(p[0], -0.1, 0.1)
        p[1] = np.clip(p[1], -0.20, 0.20)
        p[2] = np.clip(p[2], 0.0, 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[0])
        # pose.pose.position.z = float(p[2])
        pose.pose.orientation.x = 0.72
        pose.pose.orientation.y = 0.0
        pose.pose.orientation.z = 0.68
        pose.pose.orientation.w = 0.0
        pose.pose.position.x = 0.1
        pose.pose.position.y = 0.155
        pose.pose.position.z = 0.31
        goal = MoveGroup.Goal()
        req = goal.request
        req.group_name = "left_arm"
        req.allowed_planning_time = 5.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 plan_timer(self):
            if self.latest_hand is None:
                return

            # 限制 workspace
            target = self.latest_hand.copy()
            target[0] = np.clip(target[0], -0.1, 0.1)
            target[1] = np.clip(target[1], -0.2, 0.2)
            target[2] = np.clip(target[2], 0.0, 0.3)

            # 判断是否值得重新规划
            if self.last_target is not None:
                if np.linalg.norm(target - self.last_target) < self.min_delta:
                    return

            self.last_target = target.copy()
            self.send_goal(target)
            print(f"发送新目标：{target[0]:.3f}, {target[1]:.3f}, {target[2]:.3f}")

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()