#!/usr/bin/env python

# This node provides a way to define a world frame w.r.t. to a reference position.
#
# Considering that a frame position is available both:
# - w.r.t. to the world frame
# - w.r.t. to the motion capture frame
#
# It is possible to compute the transformation between the world frame
# and the motion capture world frame.
#
# This transformation is considered to be constant over the time.
#
# In the tf frame:
# world
#  mocap_world
#   ref_body
#   etc.
#  ref_body
#
# Note that the tf root frame is world.

import sys
import numpy as np

import roslib
roslib.load_manifest('evart_bridge')
import rospy

import tf
from tf.transformations import translation_from_matrix, quaternion_from_matrix
import geometry_msgs.msg
import evart_bridge.srv

trans = None
rot = None
listener = None
transformer = None

def computeTransformation(world, mocap_world, frame_world, frame_mocap):
    t = rospy.Time()

    # Retrieve the frame w.r.t. to the mocap world.
    #
    # wmMrb -> reference body position w.r.t. the motion capture world
    # wMrb -> reference body position w.r.t. the world
    #
    (transWm, rotWm) = listener.lookupTransform(frame_mocap, mocap_world, t)
    wmMrb = np.matrix(transformer.fromTranslationRotation(transWm, rotWm))
    rospy.loginfo("wmMrb: {0} {1}\n{2}".format(transWm, rotWm, wmMrb))

    (transW, rotW) = listener.lookupTransform(frame_world, world, t)
    wMrb = np.matrix(transformer.fromTranslationRotation(transW, rotW))
    rospy.loginfo("wMrb: {0} {1}\n{2}".format(transW, rotW, wMrb))

    rbMwm = np.matrix(np.linalg.inv(wmMrb))
    wMwm = np.matrix(np.linalg.inv(wMrb * rbMwm))
    trans = translation_from_matrix(wMwm)
    rot = quaternion_from_matrix(wMwm)
    rospy.loginfo("Mocap world to world tf computed: {0} {1}\n{2}".format(trans, rot, wMwm))
    return (trans, rot)

def evaluate_world_position(world, mocap_world, frame_world, frame_mocap, req):
    try:
        (trans, rot) = computeTransformation(world, mocap_world, frame_world, frame_mocap)
    except:
        return evart_bridge.srv.EvaluateWorldPositionResponse()
    return evart_bridge.srv.EvaluateWorldPositionResponse()

def evaluate_world_position_bind(world, mocap_world, frame_world, frame_mocap):
    return lambda req: evaluate_world_position(world, mocap_world,
                                               frame_world, frame_mocap, req)

if __name__ == '__main__':
    rospy.init_node('tf_world_broadcaster')
    listener = tf.TransformListener()
    transformer = tf.TransformerROS()

    # Parameters
    world = rospy.get_param('~world_frame_id', 'world')
    mocap_world = rospy.get_param('~mocap_world_frame_id', 'mocap_world')

    frame_world = rospy.get_param('~reference_body_world_frame_id')
    frame_mocap = rospy.get_param('~reference_body_mocap_frame_id')

    rospy.loginfo("world frame: {0}".format(world))
    rospy.logdebug("motion capture world frame: {0}".format(mocap_world))
    rospy.logdebug("reference body w.r.t the world frame: {0}".format(
            frame_world))
    rospy.logdebug(
        "reference body w.r.t the motion capture world frame: {0}".format(
            frame_mocap))

    # First transformation evaluation (blocking, unlike the service)
    t = rospy.Time(0)
    ok = False
    rospy.loginfo("Waiting for frames...")
    rate = rospy.Rate(10.0)
    while not ok and not rospy.is_shutdown():
        try:
            listener.waitForTransform(frame_mocap, mocap_world, t, rospy.Duration(0.1))
            listener.waitForTransform(frame_world, world, t, rospy.Duration(0.1))
            ok = True
        except tf.Exception as e:
            rospy.logdebug("error while waiting for frames: {0}".format(e))
            ok = False
            rate.sleep()
    if rospy.is_shutdown():
        sys.exit(0)

    (trans, rot) = computeTransformation(world, mocap_world,
                                         frame_world, frame_mocap)
    if rospy.is_shutdown():
        sys.exit(0)
    s = rospy.Service('evaluate_world_position',
                      evart_bridge.srv.EvaluateWorldPosition,
                      evaluate_world_position_bind(world, mocap_world,
                                                   frame_world, frame_mocap))
    tfPub = tf.TransformBroadcaster()
    rate = rospy.Rate(200)
    while not rospy.is_shutdown():
        if trans != None and rot != None:
            tfPub.sendTransform(trans, rot, rospy.Time.now(), mocap_world, world)
        rate.sleep()
