import rclpy
from rclpy.node import Node
# from control_msgs.msg import JointTrajectoryControllerState
from control_msgs.action import FollowJointTrajectory
from rclpy.action import ActionClient
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint
from sensor_msgs.msg import JointState
from std_msgs.msg import Float64MultiArray


count = 0
count1 = 0
class PositionFilter:
    """Simple per-joint filter supporting EMA and a basic 1D Kalman.

    Usage:
        f = PositionFilter(filter_type='ema', alpha=0.2)
        out = f.update(measurement)
    """
    def __init__(self, filter_type='ema', alpha=0.2, q=1e-5, r=1e-2):
        self.filter_type = filter_type
        self.alpha = float(alpha)
        # Kalman params
        self.q = float(q)  # process noise
        self.r = float(r)  # measurement noise
        self.x = None  # state estimate
        self.p = 1.0   # estimate covariance

    def reset(self):
        self.x = None
        self.p = 1.0

    def update(self, z):
        # z: measurement (float)
        if z is None:
            return self.x

        if self.filter_type == 'ema':
            if self.x is None:
                self.x = float(z)
            else:
                self.x = self.alpha * float(z) + (1.0 - self.alpha) * self.x
            return self.x

        if self.filter_type == 'kalman':
            # Predict
            if self.x is None:
                self.x = float(z)
                self.p = 1.0
                return self.x
            self.p = self.p + self.q
            # Update
            k = self.p / (self.p + self.r)
            self.x = self.x + k * (float(z) - self.x)
            self.p = (1 - k) * self.p
            return self.x

        # default: passthrough
        return float(z)

class ProgressActionClient(Node):
    
    
    def __init__(self):
        super().__init__('progress_action_client')

        # Parameters for optional position filtering
        self.declare_parameter('enable_filter', True)
        self.declare_parameter('filter_type', 'ema')
        self.declare_parameter('ema_alpha', 0.5)
        self.declare_parameter('kalman_q', 1e-5)
        self.declare_parameter('kalman_r', 1e-2)

        self.enable_filter = self.get_parameter('enable_filter').get_parameter_value().bool_value
        self.filter_type = self.get_parameter('filter_type').get_parameter_value().string_value
        self.ema_alpha = self.get_parameter('ema_alpha').get_parameter_value().double_value
        self.kalman_q = self.get_parameter('kalman_q').get_parameter_value().double_value
        self.kalman_r = self.get_parameter('kalman_r').get_parameter_value().double_value

        # Create per-joint filters (7 joints expected)
        self.num_joints = 7
        self.filters = [PositionFilter(filter_type=self.filter_type,
                                       alpha=self.ema_alpha,
                                       q=self.kalman_q,
                                       r=self.kalman_r)
                        for _ in range(self.num_joints)]

        self.subscriber_ = self.create_subscription(
            JointState, 
            "/joint_sensor_states", 
            self.listener_callback, 
            10
        )

        self.float64MultiArrayPublisher_ = self.create_publisher(Float64MultiArray, "/right_forward_position_controller/commands", 10)
        # 3-1.创建动作客户端；
        self._action_client = ActionClient(self, FollowJointTrajectory, '/right_joint_trajectory_controller/follow_joint_trajectory')
 
        # 3. 创建订阅者
        # 参数1：消息类型（如String，与发布者一致）
        # 参数2：话题名称（如"chatter"，与发布者一致）
        # 参数3：回调函数（消息到达时触发，用于处理消息）
        # 参数4：队列长度（10，缓存消息的最大数量）

        self.subscriber_  # 防止未使用变量警告
 
    # 4. 消息处理回调函数（核心）
    # 当订阅到消息时，自动调用此函数处理消息
    def listener_callback(self, msg):

        # self.send_goal([count, count1, 0.0, 0.0, 0.0, 0.0, 0.0])
        self.get_logger().info(f'Received joint states: {list(msg.name)} positions: {list(msg.position)}')
        # Apply position filtering if enabled
        arm_pos = list(msg.position)
        if self.enable_filter:
            if len(arm_pos) >= self.num_joints:
                filtered = []
                for i in range(self.num_joints):
                    try:
                        v = float(arm_pos[i])
                    except Exception:
                        v = None
                    fv = self.filters[i].update(v)
                    # If filter returned None (no prior), fall back to raw
                    filtered.append(fv if fv is not None else (arm_pos[i] if i < len(arm_pos) else 0.0))
                arm_pos_to_send = filtered
            else:
                # unexpected length: skip filtering and log
                self.get_logger().warning(f'Expected {self.num_joints} positions, got {len(arm_pos)}. Sending raw positions.')
                arm_pos_to_send = arm_pos
        else:
            arm_pos_to_send = arm_pos

        self.send_goal(arm_pos_to_send)

    def send_goal(self, arm_pos):
        arm_joint = ['openarm_right_joint1', 'openarm_right_joint2', 'openarm_right_joint3', 'openarm_right_joint4', 'openarm_right_joint5', 'openarm_right_joint6', 'openarm_right_joint7']
        arm_effort = [10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0]
        arm_zero = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        # 3-2.发送请求；
        # goal_msg = FollowJointTrajectory.Goal()
        # goal_msg.trajectory.joint_names = arm_joint
        # goal_msg.trajectory.points.append(JointTrajectoryPoint())
        # goal_msg.trajectory.points[0].positions = arm_pos
        # goal_msg.trajectory.points[0].effort = arm_effort
        # print(goal_msg)

        # self._action_client.wait_for_server()
        # self._send_goal_future = self._action_client.send_goal_async(goal_msg, feedback_callback=self.feedback_callback)
        # self._send_goal_future.add_done_callback(self.goal_response_callback)
        msg = Float64MultiArray()
        msg.data = arm_pos
        self.float64MultiArrayPublisher_.publish(msg)
        self.get_logger().info(f'Published joint commands: {arm_pos}')
 
    def goal_response_callback(self, future):
        # 3-3.处理目标发送后的反馈；
        goal_handle = future.result()
        print(goal_handle)
        if not goal_handle.accepted:
            self.get_logger().info('请求被拒绝')
            return
 
        self.get_logger().info('请求被接收，开始执行任务！')
 
        self._get_result_future = goal_handle.get_result_async()
        self._get_result_future.add_done_callback(self.get_result_callback)
 
    # 3-5.处理最终响应。
    def get_result_callback(self, future):
        return
        # result = future.result().result
        # self.get_logger().info('最终计算结果：sum = %d' % result.sum)
        # 5.释放资源。
        # rclpy.shutdown()
 
    # 3-4.处理连续反馈；
    def feedback_callback(self, feedback_msg):
        feedback = (int)(feedback_msg.feedback.progress * 100)
        self.get_logger().info('当前进度: %d%%' % feedback)

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