#!/usr/bin/env python

import numpy as np
import rospy
import geometry_msgs.msg
import tf2_ros
from tf.transformations import quaternion_slerp


def translation_to_numpy(t):
    return np.array([t.x, t.y, t.z])


def quaternion_to_numpy(q):
    return np.array([q.x, q.y, q.z, q.w])


if __name__ == '__main__':
    rospy.init_node('marker_filter')
    alpha = rospy.get_param('~alpha', 0.9)
    parent_frame_id = rospy.get_param('~parent_frame_id', 'kinect2_link')
    marker_id = rospy.get_param('~marker_id', 'marker_id0')
    marker_filtered_id = rospy.get_param(
        '~marker_filtered_id', 'marker_id0_filtered')
    rate_value = rospy.get_param('~rate_value', 125)

    tfBuffer = tf2_ros.Buffer()
    listener = tf2_ros.TransformListener(tfBuffer)
    br = tf2_ros.TransformBroadcaster()

    marker_pose = None
    marker_pose0 = None
    rate = rospy.Rate(rate_value)
    while not rospy.is_shutdown():
        marker_pose0 = marker_pose
        # Lookup the transform
        try:
            marker_pose_new = tfBuffer.lookup_transform(
                parent_frame_id, marker_id, rospy.Time())

            if not marker_pose_new is None:
                marker_pose = marker_pose_new
        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
            rospy.logwarn(e)

        if marker_pose is None:
            rate.sleep()
            continue

        # Apply running average filter to translation and rotation
        if not marker_pose0 is None:
            rotation0 = quaternion_to_numpy(marker_pose0.transform.rotation)
            rotation = quaternion_to_numpy(marker_pose.transform.rotation)
            rotation_interpolated = quaternion_slerp(
                rotation0, rotation, 1 - alpha)

            translation0 = translation_to_numpy(
                marker_pose0.transform.translation)
            translation = translation_to_numpy(
                marker_pose.transform.translation)
            translation = alpha * translation0 + (1 - alpha) * translation

            # Update pose of the marker
            marker_pose.transform.rotation.x = rotation_interpolated[0]
            marker_pose.transform.rotation.y = rotation_interpolated[1]
            marker_pose.transform.rotation.z = rotation_interpolated[2]
            marker_pose.transform.rotation.w = rotation_interpolated[3]
            marker_pose.transform.translation.x = translation[0]
            marker_pose.transform.translation.y = translation[1]
            marker_pose.transform.translation.z = translation[2]

        # Create new transform and broadcast it
        t = geometry_msgs.msg.TransformStamped()
        t.header.stamp = rospy.Time.now()
        t.header.frame_id = parent_frame_id
        t.child_frame_id = marker_filtered_id
        t.transform = marker_pose.transform
        br.sendTransform(t)
        rate.sleep()

[((11, 11, 11, 36), 'numpy.array', 'np.array', ({(11, 20, 11, 35): '[t.x, t.y, t.z]'}, {}), '([t.x, t.y, t.z])', True, 'import numpy as np\n'), ((15, 11, 15, 41), 'numpy.array', 'np.array', ({(15, 20, 15, 40): '[q.x, q.y, q.z, q.w]'}, {}), '([q.x, q.y, q.z, q.w])', True, 'import numpy as np\n'), ((19, 4, 19, 36), 'rospy.init_node', 'rospy.init_node', ({(19, 20, 19, 35): '"""marker_filter"""'}, {}), "('marker_filter')", False, 'import rospy\n'), ((20, 12, 20, 42), 'rospy.get_param', 'rospy.get_param', ({(20, 28, 20, 36): '"""~alpha"""', (20, 38, 20, 41): '0.9'}, {}), "('~alpha', 0.9)", False, 'import rospy\n'), ((21, 22, 21, 73), 'rospy.get_param', 'rospy.get_param', ({(21, 38, 21, 56): '"""~parent_frame_id"""', (21, 58, 21, 72): '"""kinect2_link"""'}, {}), "('~parent_frame_id', 'kinect2_link')", False, 'import rospy\n'), ((22, 16, 22, 59), 'rospy.get_param', 'rospy.get_param', ({(22, 32, 22, 44): '"""~marker_id"""', (22, 46, 22, 58): '"""marker_id0"""'}, {}), "('~marker_id', 'marker_id0')", False, 'import rospy\n'), ((23, 25, 24, 53), 'rospy.get_param', 'rospy.get_param', ({(24, 8, 24, 29): '"""~marker_filtered_id"""', (24, 31, 24, 52): '"""marker_id0_filtered"""'}, {}), "('~marker_filtered_id', 'marker_id0_filtered')", False, 'import rospy\n'), ((25, 17, 25, 52), 'rospy.get_param', 'rospy.get_param', ({(25, 33, 25, 46): '"""~rate_value"""', (25, 48, 25, 51): '125'}, {}), "('~rate_value', 125)", False, 'import rospy\n'), ((27, 15, 27, 31), 'tf2_ros.Buffer', 'tf2_ros.Buffer', ({}, {}), '()', False, 'import tf2_ros\n'), ((28, 15, 28, 50), 'tf2_ros.TransformListener', 'tf2_ros.TransformListener', ({(28, 41, 28, 49): 'tfBuffer'}, {}), '(tfBuffer)', False, 'import tf2_ros\n'), ((29, 9, 29, 39), 'tf2_ros.TransformBroadcaster', 'tf2_ros.TransformBroadcaster', ({}, {}), '()', False, 'import tf2_ros\n'), ((33, 11, 33, 33), 'rospy.Rate', 'rospy.Rate', ({(33, 22, 33, 32): 'rate_value'}, {}), '(rate_value)', False, 'import rospy\n'), ((34, 14, 34, 33), 'rospy.is_shutdown', 'rospy.is_shutdown', ({}, {}), '()', False, 'import rospy\n'), ((74, 25, 74, 41), 'rospy.Time.now', 'rospy.Time.now', ({}, {}), '()', False, 'import rospy\n'), ((54, 36, 55, 47), 'tf.transformations.quaternion_slerp', 'quaternion_slerp', ({(55, 16, 55, 25): 'rotation0', (55, 27, 55, 35): 'rotation', (55, 37, 55, 46): '1 - alpha'}, {}), '(rotation0, rotation, 1 - alpha)', False, 'from tf.transformations import quaternion_slerp\n'), ((39, 44, 39, 56), 'rospy.Time', 'rospy.Time', ({}, {}), '()', False, 'import rospy\n'), ((44, 12, 44, 28), 'rospy.logwarn', 'rospy.logwarn', ({(44, 26, 44, 27): 'e'}, {}), '(e)', False, 'import rospy\n')]