#!/usr/bin/env python

# Copyright (c) 2013-2015, Rethink Robotics
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
# 3. Neither the name of the Rethink Robotics nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

"""
Tool to tuck/untuck Baxter's arms to/from the shipping pose
"""
import argparse

from copy import deepcopy

import rospy

from std_msgs.msg import (
    Empty,
    Bool,
)

import baxter_interface

from baxter_core_msgs.msg import (
    CollisionAvoidanceState,
)
from baxter_interface import CHECK_VERSION


class Tuck(object):
    def __init__(self, tuck_cmd):
        self._done = False
        self._limbs = ('left', 'right')
        self._arms = {
            'left': baxter_interface.Limb('left'),
            'right': baxter_interface.Limb('right'),
            }
        self._tuck = tuck_cmd
        self._tuck_rate = rospy.Rate(20.0)  # Hz
        self._tuck_threshold = 0.2  # radians
        self._peak_angle = -1.6  # radians
        self._arm_state = {
                           'tuck': {'left': 'none', 'right': 'none'},
                           'collide': {'left': False, 'right': False},
                           'flipped': {'left': False, 'right': False}
                          }
        self._joint_moves = {
            'tuck': {
                     'left':  [-1.0, -2.07,  3.0, 2.55,  0.0, 0.01,  0.0],
                     'right':  [1.0, -2.07, -3.0, 2.55, -0.0, 0.01,  0.0]
                     },
            'untuck': {
                       'left':  [-0.08, -1.0, -1.19, 1.94,  0.67, 1.03, -0.50],
                       'right':  [0.08, -1.0,  1.19, 1.94, -0.67, 1.03,  0.50]
                       }
            }
        self._collide_lsub = rospy.Subscriber(
                             'robot/limb/left/collision_avoidance_state',
                             CollisionAvoidanceState,
                             self._update_collision, 'left')
        self._collide_rsub = rospy.Subscriber(
                             'robot/limb/right/collision_avoidance_state',
                             CollisionAvoidanceState,
                             self._update_collision, 'right')
        self._disable_pub = {
            'left': rospy.Publisher(
                 'robot/limb/left/suppress_collision_avoidance',
                 Empty, queue_size=10),
            'right': rospy.Publisher(
                 'robot/limb/right/suppress_collision_avoidance',
                 Empty, queue_size=10)
        }
        self._rs = baxter_interface.RobotEnable(CHECK_VERSION)
        self._enable_pub = rospy.Publisher('robot/set_super_enable', 
                                           Bool, queue_size=10)

    def _update_collision(self, data, limb):
        self._arm_state['collide'][limb] = len(data.collision_object) > 0
        self._check_arm_state()

    def _check_arm_state(self):
        """
        Check for goals and behind collision field.

        If s1 joint is over the peak, collision will need to be disabled
        to get the arm around the head-arm collision force-field.
        """
        diff_check = lambda a, b: abs(a - b) <= self._tuck_threshold
        for limb in self._limbs:
            angles = [self._arms[limb].joint_angle(joint)
                      for joint in self._arms[limb].joint_names()]

            # Check if in a goal position
            untuck_goal = map(diff_check, angles,
                              self._joint_moves['untuck'][limb])
            tuck_goal = map(diff_check, angles[0:2],
                            self._joint_moves['tuck'][limb][0:2])
            if all(untuck_goal):
                self._arm_state['tuck'][limb] = 'untuck'
            elif all(tuck_goal):
                self._arm_state['tuck'][limb] = 'tuck'
            else:
                self._arm_state['tuck'][limb] = 'none'

            # Check if shoulder is flipped over peak
            self._arm_state['flipped'][limb] = (
                self._arms[limb].joint_angle(limb + '_s1') <= self._peak_angle)

    def _prepare_to_tuck(self):
        # If arms are in "tucked" state, disable collision avoidance
        # before enabling robot, to avoid arm jerking from "force-field".
        head = baxter_interface.Head()
        start_disabled = not self._rs.state().enabled
        at_goal = lambda: (abs(head.pan()) <=
                        baxter_interface.settings.HEAD_PAN_ANGLE_TOLERANCE)

        rospy.loginfo("Moving head to neutral position")
        while not at_goal() and not rospy.is_shutdown():
            if start_disabled:
                [pub.publish(Empty()) for pub in self._disable_pub.values()]
            if not self._rs.state().enabled:
                self._enable_pub.publish(True)
            head.set_pan(0.0, 0.5, timeout=0)
            self._tuck_rate.sleep()

        if start_disabled:
            while self._rs.state().enabled == True and not rospy.is_shutdown():
                [pub.publish(Empty()) for pub in self._disable_pub.values()]
                self._enable_pub.publish(False)
                self._tuck_rate.sleep()

    def _move_to(self, tuck, disabled):
        if any(disabled.values()):
            [pub.publish(Empty()) for pub in self._disable_pub.values()]
        while (any(self._arm_state['tuck'][limb] != goal
                   for limb, goal in tuck.viewitems())
               and not rospy.is_shutdown()):
            if self._rs.state().enabled == False:
                self._enable_pub.publish(True)
            for limb in self._limbs:
                if disabled[limb]:
                    self._disable_pub[limb].publish(Empty())
                if limb in tuck:
                    self._arms[limb].set_joint_positions(dict(zip(
                                      self._arms[limb].joint_names(),
                                      self._joint_moves[tuck[limb]][limb])))
            self._check_arm_state()
            self._tuck_rate.sleep()

        if any(self._arm_state['collide'].values()):
            self._rs.disable()
        return

    def supervised_tuck(self):
        # Update our starting state to check if arms are tucked
        self._prepare_to_tuck()
        self._check_arm_state()
        # Tuck Arms
        if self._tuck == True:
            # If arms are already tucked, report this to user and exit.
            if all(self._arm_state['tuck'][limb] == 'tuck'
                   for limb in self._limbs):
                rospy.loginfo("Tucking: Arms already in 'Tucked' position.")
                self._done = True
                return
            else:
                rospy.loginfo("Tucking: One or more arms not Tucked.")
                any_flipped = not all(self._arm_state['flipped'].values())
                if any_flipped:
                    rospy.loginfo(
                        "Moving to neutral start position with collision %s.",
                        "on" if any_flipped else "off")
                # Move to neutral pose before tucking arms to avoid damage
                self._check_arm_state()
                actions = dict()
                disabled = {'left': True, 'right': True}
                for limb in self._limbs:
                    if not self._arm_state['flipped'][limb]:
                        actions[limb] = 'untuck'
                        disabled[limb] = False
                self._move_to(actions, disabled)

                # Disable collision and Tuck Arms
                rospy.loginfo("Tucking: Tucking with collision avoidance off.")
                actions = {'left': 'tuck', 'right': 'tuck'}
                disabled = {'left': True, 'right': True}
                self._move_to(actions, disabled)
                self._done = True
                return

        # Untuck Arms
        else:
            # If arms are tucked disable collision and untuck arms
            if any(self._arm_state['flipped'].values()):
                rospy.loginfo("Untucking: One or more arms Tucked;"
                              " Disabling Collision Avoidance and untucking.")
                self._check_arm_state()
                suppress = deepcopy(self._arm_state['flipped'])
                actions = {'left': 'untuck', 'right': 'untuck'}
                self._move_to(actions, suppress)
                self._done = True
                return
            # If arms already untucked, move to neutral location
            else:
                rospy.loginfo("Untucking: Arms already Untucked;"
                              " Moving to neutral position.")
                self._check_arm_state()
                suppress = deepcopy(self._arm_state['flipped'])
                actions = {'left': 'untuck', 'right': 'untuck'}
                self._move_to(actions, suppress)
                self._done = True
                return

    def clean_shutdown(self):
        """Handles ROS shutdown (Ctrl-C) safely."""
        if not self._done:
            rospy.logwarn('Aborting: Shutting down safely...')
        if any(self._arm_state['collide'].values()):
            while self._rs.state().enabled != False:
                [pub.publish(Empty()) for pub in self._disable_pub.values()]
                self._enable_pub.publish(False)
                self._tuck_rate.sleep()


def main():
    parser = argparse.ArgumentParser()
    tuck_group = parser.add_mutually_exclusive_group(required=True)
    tuck_group.add_argument("-t","--tuck", dest="tuck",
        action='store_true', default=False, help="tuck arms")
    tuck_group.add_argument("-u", "--untuck", dest="untuck",
        action='store_true', default=False, help="untuck arms")
    args = parser.parse_args(rospy.myargv()[1:])
    tuck = args.tuck

    rospy.loginfo("Initializing node... ")
    rospy.init_node("rsdk_tuck_arms")
    rospy.loginfo("%sucking arms" % ("T" if tuck else "Unt",))
    tucker = Tuck(tuck)
    rospy.on_shutdown(tucker.clean_shutdown)
    tucker.supervised_tuck()
    rospy.loginfo("Finished tuck")

if __name__ == "__main__":
    main()

[((248, 13, 248, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((257, 4, 257, 42), 'rospy.loginfo', 'rospy.loginfo', ({(257, 18, 257, 41): '"""Initializing node... """'}, {}), "('Initializing node... ')", False, 'import rospy\n'), ((258, 4, 258, 37), 'rospy.init_node', 'rospy.init_node', ({(258, 20, 258, 36): '"""rsdk_tuck_arms"""'}, {}), "('rsdk_tuck_arms')", False, 'import rospy\n'), ((259, 4, 259, 62), 'rospy.loginfo', 'rospy.loginfo', ({(259, 18, 259, 61): "('%sucking arms' % ('T' if tuck else 'Unt',))"}, {}), "('%sucking arms' % ('T' if tuck else 'Unt',))", False, 'import rospy\n'), ((261, 4, 261, 44), 'rospy.on_shutdown', 'rospy.on_shutdown', ({(261, 22, 261, 43): 'tucker.clean_shutdown'}, {}), '(tucker.clean_shutdown)', False, 'import rospy\n'), ((263, 4, 263, 34), 'rospy.loginfo', 'rospy.loginfo', ({(263, 18, 263, 33): '"""Finished tuck"""'}, {}), "('Finished tuck')", False, 'import rospy\n'), ((61, 26, 61, 42), 'rospy.Rate', 'rospy.Rate', ({(61, 37, 61, 41): '20.0'}, {}), '(20.0)', False, 'import rospy\n'), ((79, 29, 82, 60), 'rospy.Subscriber', 'rospy.Subscriber', ({(80, 29, 80, 72): '"""robot/limb/left/collision_avoidance_state"""', (81, 29, 81, 52): 'CollisionAvoidanceState', (82, 29, 82, 51): 'self._update_collision', (82, 53, 82, 59): '"""left"""'}, {}), "('robot/limb/left/collision_avoidance_state',\n    CollisionAvoidanceState, self._update_collision, 'left')", False, 'import rospy\n'), ((83, 29, 86, 61), 'rospy.Subscriber', 'rospy.Subscriber', ({(84, 29, 84, 73): '"""robot/limb/right/collision_avoidance_state"""', (85, 29, 85, 52): 'CollisionAvoidanceState', (86, 29, 86, 51): 'self._update_collision', (86, 53, 86, 60): '"""right"""'}, {}), "('robot/limb/right/collision_avoidance_state',\n    CollisionAvoidanceState, self._update_collision, 'right')", False, 'import rospy\n'), ((95, 19, 95, 62), 'baxter_interface.RobotEnable', 'baxter_interface.RobotEnable', ({(95, 48, 95, 61): 'CHECK_VERSION'}, {}), '(CHECK_VERSION)', False, 'import baxter_interface\n'), ((96, 27, 97, 63), 'rospy.Publisher', 'rospy.Publisher', (), '', False, 'import rospy\n'), ((134, 15, 134, 38), 'baxter_interface.Head', 'baxter_interface.Head', ({}, {}), '()', False, 'import baxter_interface\n'), ((139, 8, 139, 56), 'rospy.loginfo', 'rospy.loginfo', ({(139, 22, 139, 55): '"""Moving head to neutral position"""'}, {}), "('Moving head to neutral position')", False, 'import rospy\n'), ((57, 20, 57, 49), 'baxter_interface.Limb', 'baxter_interface.Limb', ({(57, 42, 57, 48): '"""left"""'}, {}), "('left')", False, 'import baxter_interface\n'), ((58, 21, 58, 51), 'baxter_interface.Limb', 'baxter_interface.Limb', ({(58, 43, 58, 50): '"""right"""'}, {}), "('right')", False, 'import baxter_interface\n'), ((88, 20, 90, 38), 'rospy.Publisher', 'rospy.Publisher', (), '', False, 'import rospy\n'), ((91, 21, 93, 38), 'rospy.Publisher', 'rospy.Publisher', (), '', False, 'import rospy\n'), ((239, 12, 239, 62), 'rospy.logwarn', 'rospy.logwarn', ({(239, 26, 239, 61): '"""Aborting: Shutting down safely..."""'}, {}), "('Aborting: Shutting down safely...')", False, 'import rospy\n'), ((254, 29, 254, 43), 'rospy.myargv', 'rospy.myargv', ({}, {}), '()', False, 'import rospy\n'), ((140, 36, 140, 55), 'rospy.is_shutdown', 'rospy.is_shutdown', ({}, {}), '()', False, 'import rospy\n'), ((159, 23, 159, 42), 'rospy.is_shutdown', 'rospy.is_shutdown', ({}, {}), '()', False, 'import rospy\n'), ((185, 16, 185, 76), 'rospy.loginfo', 'rospy.loginfo', ({(185, 30, 185, 75): '"""Tucking: Arms already in \'Tucked\' position."""'}, {}), '("Tucking: Arms already in \'Tucked\' position.")', False, 'import rospy\n'), ((189, 16, 189, 70), 'rospy.loginfo', 'rospy.loginfo', ({(189, 30, 189, 69): '"""Tucking: One or more arms not Tucked."""'}, {}), "('Tucking: One or more arms not Tucked.')", False, 'import rospy\n'), ((206, 16, 206, 79), 'rospy.loginfo', 'rospy.loginfo', ({(206, 30, 206, 78): '"""Tucking: Tucking with collision avoidance off."""'}, {}), "('Tucking: Tucking with collision avoidance off.')", False, 'import rospy\n'), ((217, 16, 218, 78), 'rospy.loginfo', 'rospy.loginfo', ({(217, 30, 218, 77): '"""Untucking: One or more arms Tucked; Disabling Collision Avoidance and untucking."""'}, {}), "(\n    'Untucking: One or more arms Tucked; Disabling Collision Avoidance and untucking.'\n    )", False, 'import rospy\n'), ((220, 27, 220, 63), 'copy.deepcopy', 'deepcopy', ({(220, 36, 220, 62): "self._arm_state['flipped']"}, {}), "(self._arm_state['flipped'])", False, 'from copy import deepcopy\n'), ((227, 16, 228, 61), 'rospy.loginfo', 'rospy.loginfo', ({(227, 30, 228, 60): '"""Untucking: Arms already Untucked; Moving to neutral position."""'}, {}), "('Untucking: Arms already Untucked; Moving to neutral position.')", False, 'import rospy\n'), ((230, 27, 230, 63), 'copy.deepcopy', 'deepcopy', ({(230, 36, 230, 62): "self._arm_state['flipped']"}, {}), "(self._arm_state['flipped'])", False, 'from copy import deepcopy\n'), ((149, 59, 149, 78), 'rospy.is_shutdown', 'rospy.is_shutdown', ({}, {}), '()', False, 'import rospy\n'), ((156, 25, 156, 32), 'std_msgs.msg.Empty', 'Empty', ({}, {}), '()', False, 'from std_msgs.msg import Empty, Bool\n'), ((192, 20, 194, 55), 'rospy.loginfo', 'rospy.loginfo', ({(193, 24, 193, 77): '"""Moving to neutral start position with collision %s."""', (194, 24, 194, 54): "('on' if any_flipped else 'off')"}, {}), "('Moving to neutral start position with collision %s.', 'on' if\n    any_flipped else 'off')", False, 'import rospy\n'), ((142, 29, 142, 36), 'std_msgs.msg.Empty', 'Empty', ({}, {}), '()', False, 'from std_msgs.msg import Empty, Bool\n'), ((150, 29, 150, 36), 'std_msgs.msg.Empty', 'Empty', ({}, {}), '()', False, 'from std_msgs.msg import Empty, Bool\n'), ((164, 52, 164, 59), 'std_msgs.msg.Empty', 'Empty', ({}, {}), '()', False, 'from std_msgs.msg import Empty, Bool\n'), ((242, 29, 242, 36), 'std_msgs.msg.Empty', 'Empty', ({}, {}), '()', False, 'from std_msgs.msg import Empty, Bool\n')]