#!/usr/bin/env python
import rospy
from scipy.optimize import minimize
import numpy as np
from sub8_ros_tools import wait_for_param, thread_lock, rosmsg_to_numpy
import threading
from sub8_msgs.msg import Thrust, ThrusterCmd
from geometry_msgs.msg import WrenchStamped


lock = threading.Lock()
class ThrusterMapper(object):
    def __init__(self, layout):
        '''The layout should be a dictionary of the form used in thruster_mapper.launch
        an excerpt is quoted here for posterity

        busses:
            - port: /dev/serial/by-id/usb-FTDI_FT232R_USB_UART_A403IMC9-if00-port0
              thrusters:
                FLV: {
                  node_id: 10,
                  frame_id: /base_link,
                  position: [0.1583, 0.169, 0.0142],
                  direction: [0, 0, -1]
                }
                FLL: {
                  node_id: 11,
                  frame_id: /base_link,
                  position: [0.2678, 0.2795, 0],
                  direction: [-0.866, 0.5, 0]
                }

            - port: ....
              thrusters: ....

        '''
        self.min_thrust, self.max_thrust = self.get_ranges()
        # Make thruster layout
        self.B = self.generate_B(layout)
        self.wrench_sub = rospy.Subscriber('/wrench', WrenchStamped, self.request_wrench_cb, queue_size=1)
        self.thruster_pub = rospy.Publisher('/thrust', Thrust, queue_size=1)

    @thread_lock(lock)
    def update_layout(self, srv):
        '''Respond to a thruster-out alarm'''
        raise(NotImplementedError("Layout Updates Not Implemented"))
        self.layout = '...'
        self.B = self.generate_B(self.layout)

    def get_ranges(self):
        '''Todo: Actually do something
            --> Wait for a service in the thruster driver to say something
            - We have this information in the form of calibration.json
        '''
        rospy.logwarn("Thruster mapper not querying thrusters for range data")
        return np.array([-90] * 8), np.array([90] * 8)

    def get_thruster_wrench(self, position, direction):
        '''Compute a single column of B, or the wrench created by a particular thruster'''
        forces = direction
        torques = np.cross(position, forces)
        wrench_column = np.hstack([forces, torques])
        return np.transpose(wrench_column)

    def generate_B(self, layout):
        '''Construct the control-input matrix
        Each column represents the wrench generated by a single thruster

        The single letter "B" is conventionally used to refer to a matrix which converts
         a vector of control inputs to a wrench

        Meaning where u = [thrust_1, ... thrust_n],
         B * u = [Fx, Fy, Fz, Tx, Ty, Tz]
        '''
        # Maintain an ordered list, tracking which column corresponds to which thruster
        self.thruster_name_map = []
        B = []
        for port in layout:
            for thruster_name, thruster_info in port['thrusters'].items():
                # Assemble the B matrix by columns
                self.thruster_name_map.append(thruster_name)
                wrench_column = self.get_thruster_wrench(
                    thruster_info['position'],
                    thruster_info['direction']
                )
                B.append(wrench_column)

        return np.transpose(np.array(B))

    @thread_lock(lock)
    def map(self, wrench):
        '''TODO:
            - Account for variable thrusters
        '''
        thrust_cost = np.diag([1.0] * 8)

        def objective(u):
            '''Compute a cost resembling analytical least squares
            Minimize 
                norm((B * u) - wrench) + (u.T * Q * u)
            Subject to
                min_u < u < max_u
            Where 
                Q defines the cost of firing the thrusters
                u is a vector where u[n] is the thrust output by thruster_n

            i.e., least squares with bounds
            '''
            error_cost = np.linalg.norm(self.B.dot(u) - wrench) ** 2
            effort_cost = np.transpose(u).dot(thrust_cost).dot(u)
            return error_cost + effort_cost

        def obj_jacobian(u):
            '''Compute the jacobian of the objective function

            [1] Scalar-By-Matrix derivative identities [Online]
                Available: https://en.wikipedia.org/wiki/Matrix_calculus#Scalar-by-vector_identities
            '''
            error_jacobian = 2 * self.B.T.dot(self.B.dot(u) - wrench)
            effort_jacobian = np.transpose(u).dot(2 * thrust_cost)
            return error_jacobian + effort_jacobian

        minimization = minimize(
            method='slsqp',
            fun=objective,
            jac=obj_jacobian,
            x0=(self.min_thrust + self.max_thrust) / 2,
            bounds=zip(self.min_thrust, self.max_thrust),
            tol=0.1
        )
        return minimization.x, minimization.success

    def request_wrench_cb(self, msg):
        '''Callback for requesting a wrench'''
        force = rosmsg_to_numpy(msg.wrench.force)
        torque = rosmsg_to_numpy(msg.wrench.torque)
        wrench = np.hstack([force, torque])

        u, success = self.map(wrench)
        if success:
            thrust_cmds = []
            # Assemble the list of thrust commands to send
            for name, thrust in zip(self.thruster_name_map, u):
                # > Can speed this up by avoiding appends
                thrust_cmds.append(ThrusterCmd(name=name, thrust=thrust))
            self.thruster_pub.publish(thrust_cmds)
        else:
            # I expect this to happen frequently enough that we should not raise
            rospy.logwarn("Could not achieve wrench of {}".format(wrench))


if __name__ == '__main__':
    '''
        --> Deal with ranges
        --> Deal with getting B
        --> Figure out how to communicate with the thrusters
        --> Do these things in TF
            - Decided not to do this to avoid needless traffic for a lot of static transforms
    '''
    rospy.init_node('thruster_mapper')
    busses = wait_for_param('busses')
    mapper = ThrusterMapper(busses)
    rospy.spin()
@@ -1482,18 +1482,80 @@
  '''
-Respond to
+Update the physical thruster layout.%0A        This should only be done in
  a t
@@ -1570,13 +1570,22 @@
 out 
-alarm
+event%0A        
 '''%0A
@@ -1765,16 +1765,76 @@
      '''
+Get upper and lower thrust limits for each thruster%0A        
 Todo: Ac

Clarify mapper docstrings