import rclpy
from rclpy.node import Node
from collections import deque
# 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


def map_float_range(value, in_min, in_max, out_min, out_max, clamp=False):
    """Map a float value from one range to another.

    Args:
        value: input value.
        in_min: input range minimum.
        in_max: input range maximum.
        out_min: output range minimum.
        out_max: output range maximum.
        clamp: whether to clamp output into [min(out_min, out_max), max(out_min, out_max)].

    Returns:
        Mapped float in output range.
    """
    in_span = float(in_max) - float(in_min)
    if abs(in_span) < 1e-12:
        raise ValueError('in_max and in_min must not be equal')

    out_span = float(out_max) - float(out_min)
    mapped = float(out_min) + ((float(value) - float(in_min)) / in_span) * out_span

    if clamp:
        low = min(float(out_min), float(out_max))
        high = max(float(out_min), float(out_max))
        mapped = max(low, min(high, mapped))

    return mapped


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

        self.num_joints = 8

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

        self.declare_parameter('right_joint_states_topic', '/right_joint_sensor_states')
        self.declare_parameter('left_joint_states_topic', '/left_joint_sensor_states')
        self.declare_parameter('right_arm_cmd_topic', '/right_forward_position_controller/commands')
        self.declare_parameter('right_gripper_cmd_topic', '/right_gripper_controller/commands')
        self.declare_parameter('left_arm_cmd_topic', '/left_forward_position_controller/commands')
        self.declare_parameter('left_gripper_cmd_topic', '/left_gripper_controller/commands')

        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

        self.right_joint_states_topic = self.get_parameter('right_joint_states_topic').get_parameter_value().string_value
        self.left_joint_states_topic = self.get_parameter('left_joint_states_topic').get_parameter_value().string_value
        self.right_arm_cmd_topic = self.get_parameter('right_arm_cmd_topic').get_parameter_value().string_value
        self.right_gripper_cmd_topic = self.get_parameter('right_gripper_cmd_topic').get_parameter_value().string_value
        self.left_arm_cmd_topic = self.get_parameter('left_arm_cmd_topic').get_parameter_value().string_value
        self.left_gripper_cmd_topic = self.get_parameter('left_gripper_cmd_topic').get_parameter_value().string_value

        # Right arm runtime states
        self.right_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.right_command_queue = deque(maxlen=10)

        # Left arm runtime states
        self.left_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.left_command_queue = deque(maxlen=10)

        self.right_subscriber_ = self.create_subscription(
            JointState, 
            self.right_joint_states_topic,
            self.right_listener_callback,
            10
        )

        self.left_subscriber_ = self.create_subscription(
                JointState,
                self.left_joint_states_topic,
                self.left_listener_callback,
                10,
            )

        self.right_arm_publisher = self.create_publisher(Float64MultiArray, self.right_arm_cmd_topic, 10)
        self.right_gripper_publisher = self.create_publisher(Float64MultiArray, self.right_gripper_cmd_topic, 10)
        self.left_arm_publisher = self.create_publisher(Float64MultiArray, self.left_arm_cmd_topic, 10)
        self.left_gripper_publisher = self.create_publisher(Float64MultiArray, self.left_gripper_cmd_topic, 10)

        self.create_timer(0.01, self.timer_callback)

    def right_listener_callback(self, msg):
        self._process_arm_joint_state(
            msg=msg,
            filters=self.right_filters,
            command_queue=self.right_command_queue,
        )

    def left_listener_callback(self, msg):
        self._process_arm_joint_state(
            msg=msg,
            filters=self.left_filters,
            command_queue=self.left_command_queue,
        )
    def _process_arm_joint_state(self, msg, filters, command_queue):
        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 = 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

        if len(arm_pos_to_send) < self.num_joints:
            self.get_logger().warning(
                f'Expected {self.num_joints} positions, got {len(arm_pos_to_send)}. Skip sending split commands.'
            )
            return

        # 前7个关节给机械臂控制器，第8个关节给夹爪控制器
        arm7_pos_to_send = arm_pos_to_send[:7]
        gripper_pos_to_send = arm_pos_to_send[7]

        command_queue.append((arm7_pos_to_send, gripper_pos_to_send))
    
    def timer_callback(self):
        if self.right_command_queue:
            arm7_pos_to_send, gripper_pos_to_send = self.right_command_queue.popleft()
            self.send_goal('right', arm7_pos_to_send, gripper_pos_to_send)

        if self.left_command_queue:
            arm7_pos_to_send, gripper_pos_to_send = self.left_command_queue.popleft()
            self.send_goal('left', arm7_pos_to_send, gripper_pos_to_send)

    def send_goal(self, arm_name, arm7_pos, gripper_pos):
        arm_msg = Float64MultiArray()
        arm_msg.data = arm7_pos

        gripper_msg = Float64MultiArray()
        gripper_msg.data = [map_float_range(gripper_pos, 0.0, 1.18, 0.0, 0.045)]

        if arm_name == 'left':
            self.left_arm_publisher.publish(arm_msg)
            self.left_gripper_publisher.publish(gripper_msg)
        else:
            self.right_arm_publisher.publish(arm_msg)
            self.right_gripper_publisher.publish(gripper_msg)
 

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