#!/usr/bin/env python
import rospy
#from apriltags_ros.msg import AprilTagDetectionArray
from duckietown_msgs.msg import AprilTagsWithInfos
import tf2_ros
from tf2_msgs.msg import TFMessage
import tf.transformations as tr
from geometry_msgs.msg import Transform, TransformStamped
import numpy as np
from localization import PoseAverage
from visualization_msgs.msg import Marker

# Localization Node
# Author: Teddy Ort
# Inputs: apriltags/duckietown_msgs/AprilTags - A list of april tags in a camera frame
# Outputs: pose2d/duckietown_msgs/Pose2dStamped - The estimated pose of the robot in the world frame in 2D coordinates
#          pose3d/geometry_msgs/PoseStamped - The estimated pose of the robot in the world frame in 3D coordinates

class LocalizationNode(object):
    def __init__(self):
        self.node_name = 'localization_node'

        # Constants
        self.world_frame = "world"
        self.duckiebot_frame = "duckiebot"

        self.duckiebot_lifetime = self.setupParam("~duckiebot_lifetime", 5) # The number of seconds to keep the duckiebot alive bewtween detections
        self.highlight_lifetime = self.setupParam("~highlight_lifetime", 3) # The number of seconds to keep a sign highlighted after a detection

        # Setup the publishers and subscribers
        self.sub_april = rospy.Subscriber("~apriltags", AprilTagsWithInfos, self.tag_callback)
        self.pub_tf = rospy.Publisher("/tf", TFMessage, queue_size=1, latch=True)
        self.pub_rviz = rospy.Publisher("/sign_highlights", Marker, queue_size=1, latch=True)

        # Setup the transform listener
        self.tfbuf = tf2_ros.Buffer()
        self.tfl = tf2_ros.TransformListener(self.tfbuf)

        # Use a timer to make the duckiebot disappear
        self.lifetimer = rospy.Time.now()
        self.publish_duckie_marker()

        rospy.loginfo("[%s] has started", self.node_name)

    def tag_callback(self, msg_tag):
        # Listen for the transform of the tag in the world
        avg = PoseAverage.PoseAverage()
        for tag in msg_tag.detections:
            try:
                Tt_w = self.tfbuf.lookup_transform(self.world_frame, "tag_{id}".format(id=tag.id), rospy.Time(), rospy.Duration(1))
                Mtbase_w=self.transform_to_matrix(Tt_w.transform)
                Mt_tbase = tr.concatenate_matrices(tr.translation_matrix((0,0,0.17)), tr.euler_matrix(0,0,np.pi))
                Mt_w = tr.concatenate_matrices(Mtbase_w,Mt_tbase)
                Mt_r=self.pose_to_matrix(tag.pose)
                Mr_t=np.linalg.inv(Mt_r)
                Mr_w=np.dot(Mt_w,Mr_t)
                Tr_w = self.matrix_to_transform(Mr_w)
                avg.add_pose(Tr_w)
                self.publish_sign_highlight(tag.id)
            except(tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as ex:
                rospy.logwarn("Error looking up transform for tag_%s", tag.id)
                rospy.logwarn(ex.message)

        Tr_w =  avg.get_average() # Average of the opinions

        # Broadcast the robot transform
        if Tr_w is not None:
            # Set the z translation, and x and y rotations to 0
            Tr_w.translation.z = 0
            rot = Tr_w.rotation
            rotz=tr.euler_from_quaternion((rot.x, rot.y, rot.z, rot.w))[2]
            (rot.x, rot.y, rot.z, rot.w) = tr.quaternion_from_euler(0, 0, rotz)
            T = TransformStamped()
            T.transform = Tr_w
            T.header.frame_id = self.world_frame
            T.header.stamp = rospy.Time.now()
            T.child_frame_id = self.duckiebot_frame
            self.pub_tf.publish(TFMessage([T]))
            self.lifetimer = rospy.Time.now()

    def publish_duckie_marker(self):
        # Publish a duckiebot transform far away unless the timer was reset
        rate = rospy.Rate(10)
        while not rospy.is_shutdown():
            rate.sleep()
            if rospy.Time.now() - self.lifetimer > rospy.Duration(self.duckiebot_lifetime):
                T = TransformStamped()
                T.transform.translation.z = 1000    # Throw it 1km in the air
                T.transform.rotation.w = 1
                T.header.frame_id = self.world_frame
                T.header.stamp = rospy.Time.now()
                T.child_frame_id = self.duckiebot_frame
                self.pub_tf.publish(TFMessage([T]))

    def publish_sign_highlight(self, id):
        # Publish a highlight marker on the sign that is seen by the robot
        m = Marker()
        m.header.frame_id="tag_{id}".format(id=id)
        m.header.stamp = rospy.Time.now()
        m.id=id
        m.lifetime = rospy.Duration(self.highlight_lifetime)
        m.type = Marker.CYLINDER
        p = m.pose.position
        o = m.pose.orientation
        c = m.color
        s = m.scale
        s.x, s.y, s.z = (0.1, 0.1, 0.3)
        p.z = 0.15
        c.a, c.r, c.g, c.b = (0.2, 0.9, 0.9, 0.0)
        o.w = 1
        self.pub_rviz.publish(m)

    def pose_to_matrix(self, p):
        # Return the 4x4 homogeneous matrix for a PoseStamped.msg p from the geometry_msgs
        trans = (p.pose.position.x, p.pose.position.y, p.pose.position.z)
        rot = (p.pose.orientation.x, p.pose.orientation.y, p.pose.orientation.z, p.pose.orientation.w)
        return np.dot(tr.translation_matrix(trans), tr.quaternion_matrix(rot))

    def transform_to_matrix(self, T):
        # Return the 4x4 homogeneous matrix for a TransformStamped.msg T from the geometry_msgs
        trans = (T.translation.x, T.translation.y, T.translation.z)
        rot = (T.rotation.x, T.rotation.y, T.rotation.z, T.rotation.w)
        return np.dot(tr.translation_matrix(trans), tr.quaternion_matrix(rot))

    def matrix_to_transform(self, M):
        # Return a TransformStamped.msg T from the geometry_msgs from a 4x4 homogeneous matrix
        T=Transform()
        (T.translation.x, T.translation.y, T.translation.z) = tr.translation_from_matrix(M)
        (T.rotation.x, T.rotation.y, T.rotation.z, T.rotation.w) = tr.quaternion_from_matrix(M)
        return T

    def setupParam(self, param_name, default_value):
        value = rospy.get_param(param_name, default_value)
        rospy.set_param(param_name, value)  #Write to parameter server for transparancy
        rospy.loginfo("[%s] %s = %s " % (self.node_name, param_name, value))
        return value


if __name__ == '__main__':
    rospy.init_node('localization_node', anonymous=False)
    localization_node = LocalizationNode()
    rospy.spin()

[((6381, 6434), 'rospy.init_node', 'rospy.init_node', (['"""localization_node"""'], {'anonymous': '(False)'}), "('localization_node', anonymous=False)\n", (6396, 6434), False, 'import rospy\n'), ((6482, 6494), 'rospy.spin', 'rospy.spin', ([], {}), '()\n', (6492, 6494), False, 'import rospy\n'), ((1308, 1377), 'rospy.Subscriber', 'rospy.Subscriber', (['"""~apriltags"""', 'AprilTagsWithInfos', 'self.tag_callback'], {}), "('~apriltags', AprilTagsWithInfos, self.tag_callback)\n", (1324, 1377), False, 'import rospy\n'), ((1400, 1459), 'rospy.Publisher', 'rospy.Publisher', (['"""/tf"""', 'TFMessage'], {'queue_size': '(1)', 'latch': '(True)'}), "('/tf', TFMessage, queue_size=1, latch=True)\n", (1415, 1459), False, 'import rospy\n'), ((1484, 1553), 'rospy.Publisher', 'rospy.Publisher', (['"""/sign_highlights"""', 'Marker'], {'queue_size': '(1)', 'latch': '(True)'}), "('/sign_highlights', Marker, queue_size=1, latch=True)\n", (1499, 1553), False, 'import rospy\n'), ((1615, 1631), 'tf2_ros.Buffer', 'tf2_ros.Buffer', ([], {}), '()\n', (1629, 1631), False, 'import tf2_ros\n'), ((1651, 1688), 'tf2_ros.TransformListener', 'tf2_ros.TransformListener', (['self.tfbuf'], {}), '(self.tfbuf)\n', (1676, 1688), False, 'import tf2_ros\n'), ((1769, 1785), 'rospy.Time.now', 'rospy.Time.now', ([], {}), '()\n', (1783, 1785), False, 'import rospy\n'), ((1832, 1881), 'rospy.loginfo', 'rospy.loginfo', (['"""[%s] has started"""', 'self.node_name'], {}), "('[%s] has started', self.node_name)\n", (1845, 1881), False, 'import rospy\n'), ((1993, 2018), 'localization.PoseAverage.PoseAverage', 'PoseAverage.PoseAverage', ([], {}), '()\n', (2016, 2018), False, 'from localization import PoseAverage\n'), ((3813, 3827), 'rospy.Rate', 'rospy.Rate', (['(10)'], {}), '(10)\n', (3823, 3827), False, 'import rospy\n'), ((4485, 4493), 'visualization_msgs.msg.Marker', 'Marker', ([], {}), '()\n', (4491, 4493), False, 'from visualization_msgs.msg import Marker\n'), ((4570, 4586), 'rospy.Time.now', 'rospy.Time.now', ([], {}), '()\n', (4584, 4586), False, 'import rospy\n'), ((4624, 4663), 'rospy.Duration', 'rospy.Duration', (['self.highlight_lifetime'], {}), '(self.highlight_lifetime)\n', (4638, 4663), False, 'import rospy\n'), ((5832, 5843), 'geometry_msgs.msg.Transform', 'Transform', ([], {}), '()\n', (5841, 5843), False, 'from geometry_msgs.msg import Transform, TransformStamped\n'), ((5906, 5935), 'tf.transformations.translation_from_matrix', 'tr.translation_from_matrix', (['M'], {}), '(M)\n', (5932, 5935), True, 'import tf.transformations as tr\n'), ((6003, 6031), 'tf.transformations.quaternion_from_matrix', 'tr.quaternion_from_matrix', (['M'], {}), '(M)\n', (6028, 6031), True, 'import tf.transformations as tr\n'), ((6119, 6161), 'rospy.get_param', 'rospy.get_param', (['param_name', 'default_value'], {}), '(param_name, default_value)\n', (6134, 6161), False, 'import rospy\n'), ((6170, 6204), 'rospy.set_param', 'rospy.set_param', (['param_name', 'value'], {}), '(param_name, value)\n', (6185, 6204), False, 'import rospy\n'), ((6258, 6326), 'rospy.loginfo', 'rospy.loginfo', (["('[%s] %s = %s ' % (self.node_name, param_name, value))"], {}), "('[%s] %s = %s ' % (self.node_name, param_name, value))\n", (6271, 6326), False, 'import rospy\n'), ((3340, 3376), 'tf.transformations.quaternion_from_euler', 'tr.quaternion_from_euler', (['(0)', '(0)', 'rotz'], {}), '(0, 0, rotz)\n', (3364, 3376), True, 'import tf.transformations as tr\n'), ((3393, 3411), 'geometry_msgs.msg.TransformStamped', 'TransformStamped', ([], {}), '()\n', (3409, 3411), False, 'from geometry_msgs.msg import Transform, TransformStamped\n'), ((3521, 3537), 'rospy.Time.now', 'rospy.Time.now', ([], {}), '()\n', (3535, 3537), False, 'import rospy\n'), ((3667, 3683), 'rospy.Time.now', 'rospy.Time.now', ([], {}), '()\n', (3681, 3683), False, 'import rospy\n'), ((3846, 3865), 'rospy.is_shutdown', 'rospy.is_shutdown', ([], {}), '()\n', (3863, 3865), False, 'import rospy\n'), ((5278, 5306), 'tf.transformations.translation_matrix', 'tr.translation_matrix', (['trans'], {}), '(trans)\n', (5299, 5306), True, 'import tf.transformations as tr\n'), ((5308, 5333), 'tf.transformations.quaternion_matrix', 'tr.quaternion_matrix', (['rot'], {}), '(rot)\n', (5328, 5333), True, 'import tf.transformations as tr\n'), ((5631, 5659), 'tf.transformations.translation_matrix', 'tr.translation_matrix', (['trans'], {}), '(trans)\n', (5652, 5659), True, 'import tf.transformations as tr\n'), ((5661, 5686), 'tf.transformations.quaternion_matrix', 'tr.quaternion_matrix', (['rot'], {}), '(rot)\n', (5681, 5686), True, 'import tf.transformations as tr\n'), ((2410, 2453), 'tf.transformations.concatenate_matrices', 'tr.concatenate_matrices', (['Mtbase_w', 'Mt_tbase'], {}), '(Mtbase_w, Mt_tbase)\n', (2433, 2453), True, 'import tf.transformations as tr\n'), ((2525, 2544), 'numpy.linalg.inv', 'np.linalg.inv', (['Mt_r'], {}), '(Mt_r)\n', (2538, 2544), True, 'import numpy as np\n'), ((2566, 2584), 'numpy.dot', 'np.dot', (['Mt_w', 'Mr_t'], {}), '(Mt_w, Mr_t)\n', (2572, 2584), True, 'import numpy as np\n'), ((3239, 3293), 'tf.transformations.euler_from_quaternion', 'tr.euler_from_quaternion', (['(rot.x, rot.y, rot.z, rot.w)'], {}), '((rot.x, rot.y, rot.z, rot.w))\n', (3263, 3293), True, 'import tf.transformations as tr\n'), ((3622, 3636), 'tf2_msgs.msg.TFMessage', 'TFMessage', (['[T]'], {}), '([T])\n', (3631, 3636), False, 'from tf2_msgs.msg import TFMessage\n'), ((3943, 3982), 'rospy.Duration', 'rospy.Duration', (['self.duckiebot_lifetime'], {}), '(self.duckiebot_lifetime)\n', (3957, 3982), False, 'import rospy\n'), ((4004, 4022), 'geometry_msgs.msg.TransformStamped', 'TransformStamped', ([], {}), '()\n', (4020, 4022), False, 'from geometry_msgs.msg import Transform, TransformStamped\n'), ((4230, 4246), 'rospy.Time.now', 'rospy.Time.now', ([], {}), '()\n', (4244, 4246), False, 'import rospy\n'), ((2174, 2186), 'rospy.Time', 'rospy.Time', ([], {}), '()\n', (2184, 2186), False, 'import rospy\n'), ((2188, 2205), 'rospy.Duration', 'rospy.Duration', (['(1)'], {}), '(1)\n', (2202, 2205), False, 'import rospy\n'), ((2324, 2359), 'tf.transformations.translation_matrix', 'tr.translation_matrix', (['(0, 0, 0.17)'], {}), '((0, 0, 0.17))\n', (2345, 2359), True, 'import tf.transformations as tr\n'), ((2359, 2387), 'tf.transformations.euler_matrix', 'tr.euler_matrix', (['(0)', '(0)', 'np.pi'], {}), '(0, 0, np.pi)\n', (2374, 2387), True, 'import tf.transformations as tr\n'), ((2855, 2917), 'rospy.logwarn', 'rospy.logwarn', (['"""Error looking up transform for tag_%s"""', 'tag.id'], {}), "('Error looking up transform for tag_%s', tag.id)\n", (2868, 2917), False, 'import rospy\n'), ((2934, 2959), 'rospy.logwarn', 'rospy.logwarn', (['ex.message'], {}), '(ex.message)\n', (2947, 2959), False, 'import rospy\n'), ((3907, 3923), 'rospy.Time.now', 'rospy.Time.now', ([], {}), '()\n', (3921, 3923), False, 'import rospy\n'), ((4339, 4353), 'tf2_msgs.msg.TFMessage', 'TFMessage', (['[T]'], {}), '([T])\n', (4348, 4353), False, 'from tf2_msgs.msg import TFMessage\n')]