#!/usr/bin/env python

from __future__ import division

import numpy
import scipy.optimize

import roslib
roslib.load_manifest('thruster_mapper')
import rospy
from geometry_msgs.msg import WrenchStamped
from std_msgs.msg import Header
import tf
from tf import transformations

from uf_common.orientation_helpers import xyz_array
from thruster_handling.listener import ThrusterListener


rospy.init_node('thruster_mapper')
tradeoff_distance = rospy.get_param('tradeoff_distance', 0.1) # meters
# higher tradeoff_distance -> care more about force
# lower tradeoff_distance -> care more about torque


# listen for wrenches and send out ThrusterCommand messages
tf_listener = tf.TransformListener()
thruster_listener = ThrusterListener()

def lstsq_box(A, b, x_mins, x_maxs):
    return scipy.optimize.fmin_slsqp(
        func=lambda x: numpy.linalg.norm(A.dot(x) - b)**2,
        x0=(x_mins + x_maxs)/2,
        bounds=zip(x_mins, x_maxs),
        fprime=lambda x: 2*A.T.dot(A.dot(x) - b),
        disp=0,
    )

def wrench_callback(msg):
    wrench = numpy.concatenate([xyz_array(msg.wrench.force), xyz_array(msg.wrench.torque)])
    
    # get thrusterinfos that haven't expired
    good_thrusterinfos = []
    for ti in thruster_listener.get_thrusters():
        try:
            trans, rot_q = tf_listener.lookupTransform(
                msg.header.frame_id, ti.header.frame_id, msg.header.stamp)
        except Exception, e:
            rospy.logerr("lookupTransform failed: %s", e)
            continue
        rot = transformations.quaternion_matrix(rot_q)[:3, :3]
        
        good_thrusterinfos.append((
            ti, # thrusterinfo
            trans + rot.dot(xyz_array(ti.position)), # position in wrench frame
            rot.dot(transformations.unit_vector(xyz_array(ti.direction))), # direction in wrench frame
            rot.dot(xyz_array(ti.torque_per_force)), # torque_per_force in wrench_frame
        ))
    
    if not good_thrusterinfos:
        return
    
    wrench_from_efforts = numpy.array([numpy.concatenate([
        direction, # force = direction * effort
        numpy.cross(position, direction) + torque_per_force, # torque = (position X direction + torque_per_force) * effort
    ]) for ti, position, direction, torque_per_force in good_thrusterinfos]).T
    
    # divide torques by 1/tradeoff_distance to make units match and provide
    #     a tunable parameter for the tradeoff between force and torque error
    # this section has no effect unless thrust limits are being reached
    pseudowrench_from_wrench = numpy.diag([1, 1, 1, 1/tradeoff_distance, 1/tradeoff_distance, 1/tradeoff_distance])
    pseudowrench = pseudowrench_from_wrench.dot(wrench)
    pseudowrench_from_efforts = pseudowrench_from_wrench.dot(wrench_from_efforts)
    
    # solve for efforts using box-constrained least squares
    efforts = lstsq_box(pseudowrench_from_efforts, pseudowrench,
        numpy.array([ti[0].min_force for ti in good_thrusterinfos]),
        numpy.array([ti[0].max_force for ti in good_thrusterinfos]))
    
    error = numpy.linalg.norm(pseudowrench_from_efforts.dot(efforts) - pseudowrench)
    if error > 0.01: # N
        rospy.loginfo("Thruster efforts don't exactly match desired wrench. Error: %.2f N", error)
    
    for (ti, _, _, _), effort in zip(good_thrusterinfos, efforts):
        thruster_listener.send_command(ti.id, effort)
rospy.Subscriber('wrench', WrenchStamped, wrench_callback)


rospy.spin()
