import rclpy
from rclpy.node import Node

from geometry_msgs.msg import PoseStamped
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint
from moveit_msgs.srv import GetPositionIK
from sensor_msgs.msg import JointState
import numpy as np
import time


POSITION_DEADZONE = 0.010  # 5 mm

class JointLimiter:
    def __init__(self,
                 dof,
                 max_velocity=0.5,        # rad/s
                 max_acceleration=1.0,    # rad/s^2
                 dt=0.02):                # 控制周期

        self.dof = dof
        self.max_velocity = max_velocity
        self.max_acceleration = max_acceleration
        self.dt = dt

        self.q = np.zeros(dof)
        self.dq = np.zeros(dof)
        self.initialized = False

    def reset(self, q_init):
        self.q = np.array(q_init, dtype=float)
        self.dq = np.zeros(self.dof)
        self.initialized = True

    def limit(self, q_target):
        q_target = np.array(q_target, dtype=float)

        if not self.initialized:
            self.reset(q_target)
            return q_target

        # ---------- 位置误差 ----------
        error = q_target - self.q

        # ---------- 理想速度 ----------
        dq_desired = error / self.dt

        # ---------- 速度限制 ----------
        dq_desired = np.clip(
            dq_desired,
            -self.max_velocity,
            self.max_velocity
        )

        # ---------- 加速度限制 ----------
        ddq = (dq_desired - self.dq) / self.dt

        ddq = np.clip(
            ddq,
            -self.max_acceleration,
            self.max_acceleration
        )

        # 更新速度
        self.dq = self.dq + ddq * self.dt

        # 再做一次速度保护
        self.dq = np.clip(
            self.dq,
            -self.max_velocity,
            self.max_velocity
        )

        # 更新位置
        self.q = self.q + self.dq * self.dt

        return self.q.copy()

class ArucoIKController(Node):

    def __init__(self):
        super().__init__("aruco_ik_controller")
        self.control_period = 0.02  # 20 ms
        self.joint_limiter = JointLimiter(
            dof=7,
            max_velocity=0.4,       # 建议 0.3~0.6
            max_acceleration=1.2,   # 建议 1.0~2.0
            dt=self.control_period
        )

        self.ik_ready = False
        self.prev_filtered_pos = None
        self.prev_filtered_quat = None
        self.prev_joint_angles = None
        # === 参数（根据你的 openarm 改）===
        self.arm_group = "left_arm"
        self.base_frame = "base_link"
        self.ee_link = "openarm_left_link7"

        self.joint_names = [
            "openarm_left_joint1",
            "openarm_left_joint2",
            "openarm_left_joint3",
            "openarm_left_joint4",
            "openarm_left_joint5",
            "openarm_left_joint6",
            "openarm_left_joint7",
        ]

        self.latest_joint_state = None

        self.create_subscription(
            JointState,
            "/joint_states",
            self.joint_state_cb,
            10,
        )
        self.traj_pub = self.create_publisher(
            JointTrajectory,
            "/left_joint_trajectory_controller/joint_trajectory",
            10
        )

        self.ik_client = self.create_client(
            GetPositionIK,
            "/compute_ik"
        )

        self.get_logger().info("Waiting for IK service...")
        self.ik_client.wait_for_service()

        self.create_subscription(
            PoseStamped,
            "/hand_3d",   # 已经转到 base 的 ArUco
            self.on_pose,
            10
        )

        self.timer = self.create_timer(3.0, self.check_ik_service)
    
    def joint_state_cb(self, msg: JointState):
        self.latest_joint_state = msg

    def build_seed_state(self):
        if self.latest_joint_state is None:
            return None

        seed = []
        for j in self.joint_names:
            try:
                idx = self.latest_joint_state.name.index(j)
                seed.append(self.latest_joint_state.position[idx])
            except ValueError:
                self.get_logger().error(f"Joint {j} not in joint_states")
                return None

        js = JointState()
        js.name = self.joint_names
        js.position = seed
        return js
    
    def check_ik_service(self):
        if self.ik_client.service_is_ready():
            self.get_logger().info("IK service ready!")
            self.ik_ready = True
            self.timer.cancel()

    def on_pose(self, pose: PoseStamped):
        if not self.ik_ready:
            self.get_logger().warn("IK service not ready, skip IK")
            return
        seed_state = self.build_seed_state()
        if seed_state is None:
            self.get_logger().warn("No valid seed, skip IK")
            return
        curr_pos = np.array([pose.pose.position.x, pose.pose.position.y, pose.pose.position.z])
        curr_quat = np.array([pose.pose.orientation.x, pose.pose.orientation.y, pose.pose.orientation.z, pose.pose.orientation.w])
        # self.get_logger().info(f"Received target pose: {curr_pos}, {curr_quat}")
            
        req = GetPositionIK.Request()
        req.ik_request.group_name = self.arm_group
        
        # if self.prev_filtered_pos is not None:
        #     if np.linalg.norm(curr_pos - self.prev_filtered_pos) < POSITION_DEADZONE:
        #         return  # EE 位置变化太小，跳过
            
        self.prev_filtered_pos = curr_pos
        self.prev_filtered_quat = curr_quat

        curr_pos[0] = np.clip(curr_pos[0], 0, 0.8)
        curr_pos[1] = np.clip(curr_pos[1], -0.2, 0.2) + 0.15
        curr_pos[2] = np.clip(curr_pos[2], -0.3, 0.3) + 0.3
        # pose.pose.position.x = 0.1
        # pose.pose.position.y = 0.25
        # pose.pose.position.z = 0.4
        # pose.pose.orientation.x = 1.0
        # pose.pose.orientation.y = 0.0
        # pose.pose.orientation.z = 0.0
        # pose.pose.orientation.w = 0.0

        # curr_pos = np.array([0.0, 0.0, 0.0])  # 相对于机械臂 base 的位置
        curr_quat = np.array([1.0, 0.0, 0.0, 0.0])

        

        req.ik_request.pose_stamped.pose.position.x = curr_pos[0]
        req.ik_request.pose_stamped.pose.position.y = curr_pos[1]
        req.ik_request.pose_stamped.pose.position.z = curr_pos[2]
        # req.ik_request.pose_stamped.pose.position = pose.pose.position
        req.ik_request.pose_stamped.pose.orientation.x = curr_quat[0]
        req.ik_request.pose_stamped.pose.orientation.y = curr_quat[1]
        req.ik_request.pose_stamped.pose.orientation.z = curr_quat[2]
        req.ik_request.pose_stamped.pose.orientation.w = curr_quat[3]
        req.ik_request.ik_link_name = self.ee_link
        req.ik_request.timeout.sec = 1
        req.ik_request.timeout.nanosec = 200_000_000  # 0.2s
        # req.ik_request.robot_state.joint_state.name = seed_state.name
        # req.ik_request.robot_state.joint_state.position = seed_state.position
        req.ik_request.robot_state.joint_state = seed_state
        self.get_logger().info(f"Seed joints: {seed_state.position}")
        # self.get_logger().info(f"joint names: {seed_state.name}")
        self.get_logger().info(f"Target pose: {curr_pos}, {curr_quat}")

        future = self.ik_client.call_async(req)
        future.add_done_callback(self.on_ik_result)

    def on_ik_result(self, future):
        res = future.result()
        self.get_logger().info(f"IK result code: {res.error_code.val}")
        if res.error_code.val != res.error_code.SUCCESS:
            self.get_logger().warn("IK failed")
            return
        
        jt = JointTrajectory()
        jt.joint_names = self.joint_names

        point = JointTrajectoryPoint()
        self.q_limited = self.joint_limiter.limit(res.solution.joint_state.position[:len(self.joint_names)])
        # point.positions = res.solution.joint_state.position[:len(self.joint_names)]
        point.positions = self.q_limited.tolist()
        point.time_from_start.sec = 0
        point.time_from_start.nanosec = int(self.control_period * 1e9)  # 0.3s

        jt.points.append(point)
        self.traj_pub.publish(jt)

def main():
    rclpy.init()
    node = ArucoIKController()
    rclpy.spin(node)

if __name__ == "__main__":
    main()
