#!/usr/bin/env python

import rospy
import tf
import tf2_ros
import geometry_msgs.msg
from fiducial_msgs.msg import *
import numpy as np
from math import radians as rad
from rospy import loginfo, logdebug, logwarn, logerr, logfatal


def map_cam_in_tag_frame(tf_in):
    """
    Map camera into tag frame. It is basically inverse transformation.
    """

    # copy transformation to preserve header
    tf_out = tf_in

    # parse transformation
    t = np.array(
        [[tf_in.transform.translation.x], [tf_in.transform.translation.y], [tf_in.transform.translation.z], [0]])
    q = np.array([tf_in.transform.rotation.x, tf_in.transform.rotation.y, tf_in.transform.rotation.z,
                  tf_in.transform.rotation.w])

    R = tf.transformations.quaternion_matrix(q)

    Rinv = tf.transformations.inverse_matrix(R)

    t2 = np.matmul(-Rinv, t)

    q2 = tf.transformations.quaternion_from_matrix(Rinv)

    tf_out.transform.translation.x = t2[0]
    tf_out.transform.translation.y = t2[1]
    tf_out.transform.translation.z = t2[2]
    tf_out.transform.rotation.x = q2[0]
    tf_out.transform.rotation.y = q2[1]
    tf_out.transform.rotation.z = q2[2]
    tf_out.transform.rotation.w = q2[3]

    return tf_out


def broadcast_detected_poses(msg):
    """
    Callback function for fiducial_transform topic subscriber.
    It broadcasts dynamic tf for every detected tag.
    """

    loginfo("Number of tags detected: " + str(len(msg.transforms)))

    if len(msg.transforms) > 0:
        # Init array to store all current transformations
        t = []
        # Initialize dynamic tf broadcaster
        br = tf2_ros.TransformBroadcaster()
        # Iterate all available transformations (detected tags)
        for trs in msg.transforms:
            t.append(geometry_msgs.msg.TransformStamped())
            t[-1].header.stamp = rospy.Time.now()
            t[-1].header.frame_id = "tag" + str(trs.fiducial_id)
            t[-1].child_frame_id = "tello" + str(trs.fiducial_id)
            # Both sides are the same msg type: geometry_msgs/Transform.msg
            t[-1].transform = trs.transform
            # Compute inverse transfomration to obtain camera (tello) position and orientation in tag frame
            t[-1] = map_cam_in_tag_frame(t[-1])

        # t contains all available tag-based transfomrations

        # Broadcast dynamic tf
        br.sendTransform(t)


def localization(tfBuffer):
    """
    Computed tello position based on the broadcasted transofrmations in /tf topic
    """

    trans = []
    # List of the tags to be used for localization
    used_tags = [0, 1, 2]
    for i in used_tags:
        try:
            # Transform tello# tf into map frame
            trans_tmp = tfBuffer.lookup_transform("map", "tello" + str(i), rospy.Time())
            # Do not use tf with an old timestamp (old tf remain in buffer)
            if rospy.Time.now() - trans_tmp.header.stamp < rospy.Duration(0,
                                                                          200000000):  # Duration(seconds, nanoseconds)
                trans.append(trans_tmp)
        except:
            pass

    # Compute the average position of all tellos in the map frame and broadcast final transformation
    if len(trans) > 0:
        # Init transform msg to store final tello transformation
        tello_trans = geometry_msgs.msg.TransformStamped()
        # Fill the transform
        tello_trans.header.stamp = rospy.Time.now()
        tello_trans.header.frame_id = "map"
        tello_trans.child_frame_id = "tello"
        # Translation average
        x, y, z = [], [], []
        for i in trans:
            x.append(i.transform.translation.x)
            y.append(i.transform.translation.y)
            z.append(i.transform.translation.z)
        tello_trans.transform.translation.x = sum(x) / len(x)
        tello_trans.transform.translation.y = sum(y) / len(y)
        tello_trans.transform.translation.z = sum(z) / len(z)
        # Rotation is just copied from the first transformationn in the list (just for now)
        tello_trans.transform.rotation = trans[0].transform.rotation

        # Initialize dynamic tf broadcaster
        br = tf2_ros.TransformBroadcaster()
        # Broadcast dynamic tf
        br.sendTransform(tello_trans)

        loginfo("Tello localized")
        loginfo("Number of tags used localization: " + str(len(trans)))


def main():
    """
    Main
    """

    # Init node
    rospy.init_node('localization', anonymous=False)

    # Subscribe fiducial transformation topic, the callback converts every msg to dynamic tf
    rospy.Subscriber('fiducial_transforms', FiducialTransformArray, broadcast_detected_poses)

    # Initialize tf listener
    tfBuffer = tf2_ros.Buffer()
    listener = tf2_ros.TransformListener(tfBuffer)

    rate = rospy.Rate(30)  # Hz
    # Infinite localization loop
    while not rospy.is_shutdown():
        localization(tfBuffer)
        rate.sleep()


if __name__ == '__main__':
    main()

