#!/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()

[((140, 4, 140, 57), 'rospy.init_node', 'rospy.init_node', (), '', False, 'import rospy\n'), ((142, 4, 142, 16), 'rospy.spin', 'rospy.spin', ({}, {}), '()', False, 'import rospy\n'), ((31, 25, 31, 94), 'rospy.Subscriber', 'rospy.Subscriber', ({(31, 42, 31, 54): '"""~apriltags"""', (31, 56, 31, 74): 'AprilTagsWithInfos', (31, 76, 31, 93): 'self.tag_callback'}, {}), "('~apriltags', AprilTagsWithInfos, self.tag_callback)", False, 'import rospy\n'), ((32, 22, 32, 81), 'rospy.Publisher', 'rospy.Publisher', (), '', False, 'import rospy\n'), ((33, 24, 33, 93), 'rospy.Publisher', 'rospy.Publisher', (), '', False, 'import rospy\n'), ((36, 21, 36, 37), 'tf2_ros.Buffer', 'tf2_ros.Buffer', ({}, {}), '()', False, 'import tf2_ros\n'), ((37, 19, 37, 56), 'tf2_ros.TransformListener', 'tf2_ros.TransformListener', ({(37, 45, 37, 55): 'self.tfbuf'}, {}), '(self.tfbuf)', False, 'import tf2_ros\n'), ((40, 25, 40, 41), 'rospy.Time.now', 'rospy.Time.now', ({}, {}), '()', False, 'import rospy\n'), ((43, 8, 43, 57), 'rospy.loginfo', 'rospy.loginfo', ({(43, 22, 43, 40): '"""[%s] has started"""', (43, 42, 43, 56): 'self.node_name'}, {}), "('[%s] has started', self.node_name)", False, 'import rospy\n'), ((47, 14, 47, 39), 'localization.PoseAverage.PoseAverage', 'PoseAverage.PoseAverage', ({}, {}), '()', False, 'from localization import PoseAverage\n'), ((83, 15, 83, 29), 'rospy.Rate', 'rospy.Rate', ({(83, 26, 83, 28): '10'}, {}), '(10)', False, 'import rospy\n'), ((97, 12, 97, 20), 'visualization_msgs.msg.Marker', 'Marker', ({}, {}), '()', False, 'from visualization_msgs.msg import Marker\n'), ((99, 25, 99, 41), 'rospy.Time.now', 'rospy.Time.now', ({}, {}), '()', False, 'import rospy\n'), ((101, 21, 101, 60), 'rospy.Duration', 'rospy.Duration', ({(101, 36, 101, 59): 'self.highlight_lifetime'}, {}), '(self.highlight_lifetime)', False, 'import rospy\n'), ((127, 10, 127, 21), 'geometry_msgs.msg.Transform', 'Transform', ({}, {}), '()', False, 'from geometry_msgs.msg import Transform, TransformStamped\n'), ((128, 62, 128, 91), 'tf.transformations.translation_from_matrix', 'tr.translation_from_matrix', ({(128, 89, 128, 90): 'M'}, {}), '(M)', True, 'import tf.transformations as tr\n'), ((129, 67, 129, 95), 'tf.transformations.quaternion_from_matrix', 'tr.quaternion_from_matrix', ({(129, 93, 129, 94): 'M'}, {}), '(M)', True, 'import tf.transformations as tr\n'), ((133, 16, 133, 58), 'rospy.get_param', 'rospy.get_param', ({(133, 32, 133, 42): 'param_name', (133, 44, 133, 57): 'default_value'}, {}), '(param_name, default_value)', False, 'import rospy\n'), ((134, 8, 134, 42), 'rospy.set_param', 'rospy.set_param', ({(134, 24, 134, 34): 'param_name', (134, 36, 134, 41): 'value'}, {}), '(param_name, value)', False, 'import rospy\n'), ((135, 8, 135, 76), 'rospy.loginfo', 'rospy.loginfo', ({(135, 22, 135, 75): "('[%s] %s = %s ' % (self.node_name, param_name, value))"}, {}), "('[%s] %s = %s ' % (self.node_name, param_name, value))", False, 'import rospy\n'), ((72, 43, 72, 79), 'tf.transformations.quaternion_from_euler', 'tr.quaternion_from_euler', ({(72, 68, 72, 69): '0', (72, 71, 72, 72): '0', (72, 74, 72, 78): 'rotz'}, {}), '(0, 0, rotz)', True, 'import tf.transformations as tr\n'), ((73, 16, 73, 34), 'geometry_msgs.msg.TransformStamped', 'TransformStamped', ({}, {}), '()', False, 'from geometry_msgs.msg import Transform, TransformStamped\n'), ((76, 29, 76, 45), 'rospy.Time.now', 'rospy.Time.now', ({}, {}), '()', False, 'import rospy\n'), ((79, 29, 79, 45), 'rospy.Time.now', 'rospy.Time.now', ({}, {}), '()', False, 'import rospy\n'), ((84, 18, 84, 37), 'rospy.is_shutdown', 'rospy.is_shutdown', ({}, {}), '()', False, 'import rospy\n'), ((117, 22, 117, 50), 'tf.transformations.translation_matrix', 'tr.translation_matrix', ({(117, 44, 117, 49): 'trans'}, {}), '(trans)', True, 'import tf.transformations as tr\n'), ((117, 52, 117, 77), 'tf.transformations.quaternion_matrix', 'tr.quaternion_matrix', ({(117, 73, 117, 76): 'rot'}, {}), '(rot)', True, 'import tf.transformations as tr\n'), ((123, 22, 123, 50), 'tf.transformations.translation_matrix', 'tr.translation_matrix', ({(123, 44, 123, 49): 'trans'}, {}), '(trans)', True, 'import tf.transformations as tr\n'), ((123, 52, 123, 77), 'tf.transformations.quaternion_matrix', 'tr.quaternion_matrix', ({(123, 73, 123, 76): 'rot'}, {}), '(rot)', True, 'import tf.transformations as tr\n'), ((53, 23, 53, 65), 'tf.transformations.concatenate_matrices', 'tr.concatenate_matrices', ({(53, 47, 53, 55): 'Mtbase_w', (53, 56, 53, 64): 'Mt_tbase'}, {}), '(Mtbase_w, Mt_tbase)', True, 'import tf.transformations as tr\n'), ((55, 21, 55, 40), 'numpy.linalg.inv', 'np.linalg.inv', ({(55, 35, 55, 39): 'Mt_r'}, {}), '(Mt_r)', True, 'import numpy as np\n'), ((56, 21, 56, 38), 'numpy.dot', 'np.dot', ({(56, 28, 56, 32): 'Mt_w', (56, 33, 56, 37): 'Mr_t'}, {}), '(Mt_w, Mr_t)', True, 'import numpy as np\n'), ((71, 17, 71, 71), 'tf.transformations.euler_from_quaternion', 'tr.euler_from_quaternion', ({(71, 42, 71, 70): '(rot.x, rot.y, rot.z, rot.w)'}, {}), '((rot.x, rot.y, rot.z, rot.w))', True, 'import tf.transformations as tr\n'), ((78, 32, 78, 46), 'tf2_msgs.msg.TFMessage', 'TFMessage', ({(78, 42, 78, 45): '[T]'}, {}), '([T])', False, 'from tf2_msgs.msg import TFMessage\n'), ((86, 51, 86, 90), 'rospy.Duration', 'rospy.Duration', ({(86, 66, 86, 89): 'self.duckiebot_lifetime'}, {}), '(self.duckiebot_lifetime)', False, 'import rospy\n'), ((87, 20, 87, 38), 'geometry_msgs.msg.TransformStamped', 'TransformStamped', ({}, {}), '()', False, 'from geometry_msgs.msg import Transform, TransformStamped\n'), ((91, 33, 91, 49), 'rospy.Time.now', 'rospy.Time.now', ({}, {}), '()', False, 'import rospy\n'), ((50, 99, 50, 111), 'rospy.Time', 'rospy.Time', ({}, {}), '()', False, 'import rospy\n'), ((50, 113, 50, 130), 'rospy.Duration', 'rospy.Duration', ({(50, 128, 50, 129): '1'}, {}), '(1)', False, 'import rospy\n'), ((52, 51, 52, 84), 'tf.transformations.translation_matrix', 'tr.translation_matrix', ({(52, 73, 52, 83): '(0, 0, 0.17)'}, {}), '((0, 0, 0.17))', True, 'import tf.transformations as tr\n'), ((52, 86, 52, 112), 'tf.transformations.euler_matrix', 'tr.euler_matrix', ({(52, 102, 52, 103): '0', (52, 104, 52, 105): '0', (52, 106, 52, 111): 'np.pi'}, {}), '(0, 0, np.pi)', True, 'import tf.transformations as tr\n'), ((61, 16, 61, 78), 'rospy.logwarn', 'rospy.logwarn', ({(61, 30, 61, 69): '"""Error looking up transform for tag_%s"""', (61, 71, 61, 77): 'tag.id'}, {}), "('Error looking up transform for tag_%s', tag.id)", False, 'import rospy\n'), ((62, 16, 62, 41), 'rospy.logwarn', 'rospy.logwarn', ({(62, 30, 62, 40): 'ex.message'}, {}), '(ex.message)', False, 'import rospy\n'), ((86, 15, 86, 31), 'rospy.Time.now', 'rospy.Time.now', ({}, {}), '()', False, 'import rospy\n'), ((93, 36, 93, 50), 'tf2_msgs.msg.TFMessage', 'TFMessage', ({(93, 46, 93, 49): '[T]'}, {}), '([T])', False, 'from tf2_msgs.msg import TFMessage\n')]