<reponame>stanFurrer/Multimodal-solution-for-grasp-stability-prediction<gh_stars>0
#!/usr/bin/env python
#
# Copyright (C) 2021 Learning Algorithms and Systems Laboratory, EPFL, Switzerland
# Authors:
# 	<NAME> (<EMAIL>)
# 
# Website: lasa.epfl.ch
#
# This file is part of iiwa_gazebo.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#
import rospy
import numpy as np
import argparse
import time
import tf2_ros
from tf2_geometry_msgs import PoseStamped
from std_msgs.msg import String
from geometry_msgs.msg import Pose

# ----- Argument parser -----
parser = argparse.ArgumentParser(description='User can decide to pass ee_pose_d (default) or to pass a ee_vel_d')
parser.add_argument('--velocity', 
                    type=bool, default=False,
                    help='Set the ee_vel_d mode (default False)')
parser.add_argument('--v',
                    type=float, default=0,
                    help='Set desired velocity v m/s')
parser.add_argument('--nb_sub_rot',
                    type=int, default=1,
                    help='Number of integer sub_rotations e.g 4 sub_rot of 45 deg -> 180 deg')
args = parser.parse_args()

v_max = 1.0
v_min = 0
if args.v != 0:
    if args.v > v_max:
        args.v = v_max
    elif args.v < v_min:
        args.v = v_min
    print('desired velocity is {}'.format(args.v))
v_des = args.v

if args.nb_sub_rot < 1:
    args.nb_sub_rot = 1

# ----- Helper functions for quaternions-----
def euler_to_quaternion(roll, pitch, yaw):

    qx = np.sin(roll/2) * np.cos(pitch/2) * np.cos(yaw/2) - np.cos(roll/2) * np.sin(pitch/2) * np.sin(yaw/2)
    qy = np.cos(roll/2) * np.sin(pitch/2) * np.cos(yaw/2) + np.sin(roll/2) * np.cos(pitch/2) * np.sin(yaw/2)
    qz = np.cos(roll/2) * np.cos(pitch/2) * np.sin(yaw/2) - np.sin(roll/2) * np.sin(pitch/2) * np.cos(yaw/2)
    qw = np.cos(roll/2) * np.cos(pitch/2) * np.cos(yaw/2) + np.sin(roll/2) * np.sin(pitch/2) * np.sin(yaw/2)

    return [qx, qy, qz, qw]

def normalize_quat(q):
    """ Normalise the quaternion

    Args:
        q (Quaternion)
    Return:
        q (Quaternion): normalized quaternion
    """
    norm = np.linalg.norm([q.x, q.y, q.z, q.w])
    q.x = q.x / norm
    q.y = q.y / norm
    q.z = q.z / norm
    q.w = q.w / norm
    return q

def angle_between_quaternion(p1, p2):
    """Return angle between two quaternions

    Args:
        p1 (Quaternion): First Quaternion
        p2 (Quaternion): Second Quaternion

    Returns:
        float: Angle between the quaternions [rad]
    """
    q1 = np.array([p1.x, p1.y, p1.z, p1.w])
    q2 = np.array([p2.x, p2.y, p2.z, p2.w])

    # Inv of q2
    q2_conj = q2 * [-1, -1, -1, 1]
    q2_inv = q2_conj / np.linalg.norm(q2)**2

    # norm [ vector component of q1 * q2^(-1) ] = sin(\theta / 2)
    temp = q1[3] * q2_inv[:3] + q2_inv[3] * q1[:3] + np.cross(q1[:3], q2_inv[:3])
    return 2 * np.arcsin(np.linalg.norm(temp))


# ----- Movements (predefined) -----
lat_displ = 0.1 # m 
ang_rot = np.pi/1.5 # radians
# Desired rotation in the iiwa_link_7 reference frame
q_rot_cw = euler_to_quaternion(0.0, 0.0, ang_rot)
q_rot_ccw = euler_to_quaternion(0.0, 0.0, -ang_rot) 
MOVEMENTS = {
    "home": [0.8, 0.0, 0.4, 0.0, 0.7071068, 0.0, 0.7071068], # Reset default pose
    "right": [0.0, -lat_displ, 0.0, 0.0, 0.0, 0.0, 0.0],  # Move right
    "left": [0.0, lat_displ, 0.0, 0.0, 0.0, 0.0, 0.0],  # Move left
    "down": [0.0, 0.0, -lat_displ, 0.0, 0.0, 0.0, 0.0],  # Move downwards
    "up": [0.0, 0.0, lat_displ, 0.0, 0.0, 0.0, 0.0],  # Move upwards
    "clockwise": [0.0, 0.0, 0.0, q_rot_cw[0], q_rot_cw[1], q_rot_cw[2], q_rot_cw[3]],  # Rotate clockwise around x axis
    "cclockwise": [0.0, 0.0, 0.0, q_rot_ccw[0], q_rot_ccw[1], q_rot_ccw[2], q_rot_ccw[3]],  # Rotate counter-clockwise 
    "vertical": "vertical", # Sequence downward upward
    "plane": "plane", # Sequence right left
    "fullcw": "fullcw", 
    "fullccw": "fullccw",
}

# ----- Global Variables -----
TARGET = Pose()
MOVEMENT = MOVEMENTS["home"]
EE_POSE = Pose()
keyboard_cmd = False
tf_buffer = None
is_home = False
is_rotation_cw = False
is_rotation_ccw = False
is_sequence = False
end_cw = False
end_ccw = False
frame_id_ee = "iiwa_link_7" 
delta_t = 0.0

# ----- Callbacks -----
def keyboard_callback(data):
    """Callback for keyboard commands

    Args:
        data (String): ROS `String` message
    """
    global MOVEMENT
    global keyboard_cmd
    global is_home
    global is_rotation_cw
    global is_rotation_ccw
    global is_sequence
    # Check if known target, otherwise keep default pose
    MOVEMENT = MOVEMENTS.get(data.data, MOVEMENTS["home"])
    if data.data == "home":
        is_home = True
    if (data.data == "clockwise") or (data.data == "fullcw"):
        is_rotation_cw = True
    if (data.data == "cclockwise") or (data.data == "fullccw"):
        is_rotation_ccw = True
    if (data.data == "vertical") or (data.data == "plane") or (data.data == "fullcw") or (data.data == "fullccw"):
        is_sequence = True
    keyboard_cmd = True

def ee_pose_callback(data):
    """Callback for end effector pose

    Args:
        data (Pose): ROS `Pose` message
    """
    global EE_POSE
    EE_POSE = data

# ----- Helper functions for pose and target computation-----
def to_pose():
    """Transform corresponding MOVEMENT in Pose object
    Return: 
        pose (Pose)
    """
    pose = Pose()
    pose.position.x = MOVEMENT[0]
    pose.position.y = MOVEMENT[1]
    pose.position.z = MOVEMENT[2]
    pose.orientation.x = MOVEMENT[3]
    pose.orientation.y = MOVEMENT[4]
    pose.orientation.z = MOVEMENT[5]
    pose.orientation.w = MOVEMENT[6]
    return pose

def sum_poses(pose1, pose2):
    """Sum the position and orientation fields between two poses

    Args:
        pose1 (Pose)
        pose2 (Pose)
    Return:
        sum_p (Pose)
    """
    sum_p = Pose()
    sum_p.position.x = pose1.position.x + pose2.position.x
    sum_p.position.y = pose1.position.y + pose2.position.y
    sum_p.position.z = pose1.position.z + pose2.position.z
    sum_p.orientation.x = pose1.orientation.x + pose2.orientation.x
    sum_p.orientation.y = pose1.orientation.y + pose2.orientation.y
    sum_p.orientation.z = pose1.orientation.z + pose2.orientation.z
    sum_p.orientation.w = pose1.orientation.w + pose2.orientation.w
    if is_rotation_cw or is_rotation_ccw or is_home:
        sum_p.orientation.x = pose2.orientation.x
        sum_p.orientation.y = pose2.orientation.y 
        sum_p.orientation.z = pose2.orientation.z 
        sum_p.orientation.w = pose2.orientation.w 
    
    return check_feasible(sum_p)

def to_world_frame(pose_mvt):
    """Use tf transform to express the desired rotation in the world frame

    Args:
        pose_mvt (Pose):pose associated to the desired movement in the ee frame
    Return:
        rel_target_pose (Pose): relative desired pose in the world frame
    """
    # Apply tf only if some rotation is involved 
    if (is_rotation_cw) or (is_rotation_ccw):
        try:
            pose_stamped = PoseStamped()
            pose_stamped.pose = pose_mvt
            pose_stamped.header.frame_id = frame_id_ee
            # Hack to ensure that the target transformation is always in future
            pose_stamped.header.stamp = rospy.Time.from_sec(rospy.Time.now().to_sec() + 0.5)
            output_pose_stamped = tf_buffer.transform(pose_stamped, base_link, timeout=rospy.Duration(1.0))
        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException):
            raise
        rel_target_pose = output_pose_stamped.pose
        rel_target_pose.position = pose_mvt.position # The desired position was already expressed in the world frame
    else:
        rel_target_pose = pose_mvt
   
    return  rel_target_pose

def check_feasible(p):
    """Check if boundary conditions are respected, constraint ee rotation to physical limits

    Args: 
        sum_p (Pose): computed sum of two poses
    Return:
        sum_p (Pose): saturate to limits if necessary 
    """
    global is_rotation_cw
    global is_rotation_ccw
    global end_cw
    global end_ccw
    
    if (is_rotation_cw):
        if (not end_ccw and np.abs(p.orientation.w + p.orientation.y) < 0.3) or end_cw:
            # Saturate all orientations to prevent overflow and keep quaternion normalized
            p.orientation.x = 0.7071068
            p.orientation.y = 0.0
            p.orientation.z = 0.7071068
            p.orientation.w = 0.0
            rospy.loginfo("cw saturation")
            end_cw = True
        end_ccw = False
        is_rotation_cw = False
    if (is_rotation_ccw):
        if (not end_cw and np.abs(p.orientation.w + p.orientation.y) < 0.3) or end_ccw:
            p.orientation.x = 0.7071068
            p.orientation.y = 0.0
            p.orientation.z = 0.7071068
            p.orientation.w = 0.0
            rospy.loginfo("ccw saturation")
            end_ccw = True
        end_cw = False
        is_rotation_ccw = False
    return p 

def get_error(p1, p2):
    """Return error between two poses

    Args:
        p1 (Pose): First Pose
        p2 (Pose): Second Pose

    Returns:
        Tuple[float, float]: Tuple of position and orientation error
            Position error: L2 distance between the poses
            Orientation error: Angle between the poses [rad]
    """
    pos_err = np.array([
        np.linalg.norm([p1.position.x - p2.position.x]),
        np.linalg.norm([p1.position.y - p2.position.y]),
        np.linalg.norm([p1.position.z - p2.position.z]),
    ])
    orient_err = np.abs(angle_between_quaternion(p1.orientation, p2.orientation))

    return pos_err, orient_err


# ----- Trajectories definition -----
def pub_target(rel_mvt):
    global TARGET
    TARGET = sum_poses(EE_POSE, rel_mvt)
    TARGET.orientation = normalize_quat(TARGET.orientation)
    publisher.publish(TARGET)

def reach_lin_target(nb_axis, radius_norm = 0.01):
    while not rospy.is_shutdown():
        err = get_error(EE_POSE, TARGET)
        
        if(np.linalg.norm(err[0][nb_axis]) < radius_norm): #define norm to change waypoint
            break

def reach_ang_target(angular_norm = 0.3):
    while not rospy.is_shutdown():
            err = get_error(EE_POSE, TARGET)
            if(np.linalg.norm(err[1]) < angular_norm):
                break

def vertical_mvt(rep, lat_displ, dt):
    global MOVEMENT
    mvt = Pose()
    mvt.position.x = 0
    mvt.position.y = 0
    mvt.position.z = 0
    mvt.orientation.x = EE_POSE.orientation.x
    mvt.orientation.y = EE_POSE.orientation.y
    mvt.orientation.z = EE_POSE.orientation.z
    mvt.orientation.w = EE_POSE.orientation.w
    for i in range(rep):
        if v_des != 0:
           
            mvt.position.z = -v_des
            TARGET = mvt
            publisher_v.publish(TARGET)
            t_init = time.time()
            while (time.time() - t_init < dt):
                pass
        
            mvt.position.z = v_des
           
            TARGET = mvt
            publisher_v.publish(TARGET)
            t_init = time.time()
            while (time.time() - t_init < dt + 0.03):
                pass
                
        else:
            
            rel_mvt = Pose()
            rel_mvt.position.z = -lat_displ
            pub_target(rel_mvt)
            reach_lin_target(2)
            rel_mvt.position.z = lat_displ
            pub_target(rel_mvt)
            reach_lin_target(2)
            
    #Reset position
    MOVEMENT = MOVEMENTS["home"]
    TARGET = to_pose()
    publisher.publish(TARGET)
  
        
def horizontal_mvt(rep, lat_displ, dt):
    global MOVEMENT
    mvt = Pose()
    mvt.position.x = 0
    mvt.position.y = 0
    mvt.position.z = 0
    mvt.orientation.x = EE_POSE.orientation.x
    mvt.orientation.y = EE_POSE.orientation.y
    mvt.orientation.z = EE_POSE.orientation.z
    mvt.orientation.w = EE_POSE.orientation.w
    for i in range(rep):
        if v_des != 0:
            
            mvt.position.y = -v_des
        
            TARGET = mvt
            publisher_v.publish(TARGET)
            t_init = time.time()
            while (time.time() - t_init < dt):
                pass

            mvt.position.y = v_des
            
            TARGET = mvt
            publisher_v.publish(TARGET)
            t_init = time.time()
            while (time.time() - t_init < dt):
                pass
            
        else:
            rel_mvt = Pose()
            rel_mvt.position.y = -lat_displ
            pub_target(rel_mvt)
            reach_lin_target(1)
            rel_mvt.position.y = lat_displ
            pub_target(rel_mvt)
            reach_lin_target(1)

    #Reset position
    MOVEMENT = MOVEMENTS["home"]
    TARGET = to_pose()
    publisher.publish(TARGET)
        

def full_cw(ang_rot, rep = 2):

    global MOVEMENT
    global is_rotation_cw
    global is_rotation_ccw

    for i in range(rep):
        is_rotation_cw = True #otherwise not computing tf
        rel_mvt = Pose()
        q_rot_cw = euler_to_quaternion(0.0, 0.0, ang_rot) 
        MOVEMENT = [0.0, 0.0, 0.0, q_rot_cw[0], q_rot_cw[1], q_rot_cw[2], q_rot_cw[3]]
        pose_mvt = to_pose()
        rel_mvt = to_world_frame(pose_mvt)
        pub_target(rel_mvt)
        reach_ang_target()

        is_rotation_ccw =True #otherwise not computing tf
        q_rot_ccw = euler_to_quaternion(0.0, 0.0, -ang_rot) 
        MOVEMENT = [0.0, 0.0, 0.0, q_rot_ccw[0], q_rot_ccw[1], q_rot_ccw[2], q_rot_ccw[3]]
        pose_mvt = to_pose()
        rel_mvt = to_world_frame(pose_mvt)
        pub_target(rel_mvt)
        reach_ang_target()
        
    #Reset position
    MOVEMENT = MOVEMENTS["home"]
    TARGET = to_pose()
    publisher.publish(TARGET)
  
        
def full_ccw(ang_rot, rep = 2):
    global MOVEMENT
    global is_rotation_cw
    global is_rotation_ccw

    for i in range(rep):
        is_rotation_ccw = True
        rel_mvt = Pose()
        q_rot_ccw = euler_to_quaternion(0.0, 0.0, -ang_rot) 
        MOVEMENT = [0.0, 0.0, 0.0, q_rot_ccw[0], q_rot_ccw[1], q_rot_ccw[2], q_rot_ccw[3]]
        pose_mvt = to_pose()
        rel_mvt = to_world_frame(pose_mvt)
        pub_target(rel_mvt)
        reach_ang_target()
       
        is_rotation_cw = True #otherwise not computing tf
        q_rot_cw = euler_to_quaternion(0.0, 0.0, ang_rot) 
        MOVEMENT = [0.0, 0.0, 0.0, q_rot_cw[0], q_rot_cw[1], q_rot_cw[2], q_rot_cw[3]]
        pose_mvt = to_pose()
        rel_mvt = to_world_frame(pose_mvt)
        pub_target(rel_mvt)
        reach_ang_target()

    #Reset position
    MOVEMENT = MOVEMENTS["home"]
    TARGET = to_pose()
    publisher.publish(TARGET)

def run_sequence():
    
    global delta_t
    if v_des != 0:
        delta_t = 0.3
    if MOVEMENT == "vertical":
        vertical_mvt(rep, lat_displ, delta_t)
    if MOVEMENT == "plane":
        horizontal_mvt(rep, lat_displ, delta_t)
    if MOVEMENT == "fullcw":
        full_cw(ang_rot)
    if MOVEMENT == "fullccw":
        full_ccw(ang_rot)

# ----- Main Script -----
if __name__ == '__main__':
    rospy.init_node('gazebo_iiwa_keyboard_cmd')
    rate = rospy.Rate(30)
    base_link = rospy.get_param("~base_link", "world")

    # Subscribers
    tf_buffer = tf2_ros.Buffer()
    listener = tf2_ros.TransformListener(tf_buffer)
    rospy.Subscriber("iiwa/ee_info/Pose", Pose, ee_pose_callback)
    rospy.Subscriber("iiwa/lib_cmd", String, keyboard_callback, queue_size=1)

    # Publishers
    publisher = rospy.Publisher("passive_control/pos_quat", Pose, queue_size=1)

    #Publisher to ee_vel_d
    publisher_v = rospy.Publisher("passive_control/vel_quat", Pose, queue_size=1)

    # Main loop
    while not rospy.is_shutdown():
        rep = 2
        if args.velocity and keyboard_cmd:
            if is_home:
                TARGET = to_pose()
                is_home = False
                TARGET.orientation = normalize_quat(TARGET.orientation)
                publisher.publish(TARGET)
            # Rotation
            elif is_rotation_cw or is_rotation_ccw:
                ang_rot = ang_rot/args.nb_sub_rot
                run_sequence()  
            
            else: #Linear movement
                v_des = args.v
                if v_des == 0:
                    rospy.loginfo('ATTENTION DESIRED VELOCITY IS 0')
                run_sequence()  
                
            is_sequence = False
            keyboard_cmd = False

        else:

            if(is_sequence): 
                run_sequence()  
                is_sequence = False
                keyboard_cmd = False
            elif(keyboard_cmd):
                if is_home:
                    TARGET = to_pose()
                    is_home = False
                else:
                    pose_mvt = to_pose()
                    rel_target_pose = to_world_frame(pose_mvt)
                    TARGET = sum_poses(EE_POSE, rel_target_pose)
                TARGET.orientation = normalize_quat(TARGET.orientation)
                publisher.publish(TARGET)
                keyboard_cmd = False
                rospy.loginfo(TARGET) # print for sanity check
           
        rate.sleep()
    

#!/usr/bin/env python
#
# Copyright (C) 2021 Learning Algorithms and Systems Laboratory, EPFL, Switzerland
# Authors:
# 	<NAME> (<EMAIL>)
# 
# Website: lasa.epfl.ch
#
# This file is part of iiwa_gazebo.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#
import rospy
import numpy as np
import argparse
import time
import tf2_ros
from tf2_geometry_msgs import PoseStamped
from std_msgs.msg import String
from geometry_msgs.msg import Pose

# ----- Argument parser -----
parser = argparse.ArgumentParser(description='User can decide to pass ee_pose_d (default) or to pass a ee_vel_d')
parser.add_argument('--velocity', 
                    type=bool, default=False,
                    help='Set the ee_vel_d mode (default False)')
parser.add_argument('--v',
                    type=float, default=0,
                    help='Set desired velocity v m/s')
parser.add_argument('--nb_sub_rot',
                    type=int, default=1,
                    help='Number of integer sub_rotations e.g 4 sub_rot of 45 deg -> 180 deg')
args = parser.parse_args()

v_max = 1.0
v_min = 0
if args.v != 0:
    if args.v > v_max:
        args.v = v_max
    elif args.v < v_min:
        args.v = v_min
    print('desired velocity is {}'.format(args.v))
v_des = args.v

if args.nb_sub_rot < 1:
    args.nb_sub_rot = 1

# ----- Helper functions for quaternions-----
def euler_to_quaternion(roll, pitch, yaw):

    qx = np.sin(roll/2) * np.cos(pitch/2) * np.cos(yaw/2) - np.cos(roll/2) * np.sin(pitch/2) * np.sin(yaw/2)
    qy = np.cos(roll/2) * np.sin(pitch/2) * np.cos(yaw/2) + np.sin(roll/2) * np.cos(pitch/2) * np.sin(yaw/2)
    qz = np.cos(roll/2) * np.cos(pitch/2) * np.sin(yaw/2) - np.sin(roll/2) * np.sin(pitch/2) * np.cos(yaw/2)
    qw = np.cos(roll/2) * np.cos(pitch/2) * np.cos(yaw/2) + np.sin(roll/2) * np.sin(pitch/2) * np.sin(yaw/2)

    return [qx, qy, qz, qw]

def normalize_quat(q):
    """ Normalise the quaternion

    Args:
        q (Quaternion)
    Return:
        q (Quaternion): normalized quaternion
    """
    norm = np.linalg.norm([q.x, q.y, q.z, q.w])
    q.x = q.x / norm
    q.y = q.y / norm
    q.z = q.z / norm
    q.w = q.w / norm
    return q

def angle_between_quaternion(p1, p2):
    """Return angle between two quaternions

    Args:
        p1 (Quaternion): First Quaternion
        p2 (Quaternion): Second Quaternion

    Returns:
        float: Angle between the quaternions [rad]
    """
    q1 = np.array([p1.x, p1.y, p1.z, p1.w])
    q2 = np.array([p2.x, p2.y, p2.z, p2.w])

    # Inv of q2
    q2_conj = q2 * [-1, -1, -1, 1]
    q2_inv = q2_conj / np.linalg.norm(q2)**2

    # norm [ vector component of q1 * q2^(-1) ] = sin(\theta / 2)
    temp = q1[3] * q2_inv[:3] + q2_inv[3] * q1[:3] + np.cross(q1[:3], q2_inv[:3])
    return 2 * np.arcsin(np.linalg.norm(temp))


# ----- Movements (predefined) -----
lat_displ = 0.1 # m 
ang_rot = np.pi/1.5 # radians
# Desired rotation in the iiwa_link_7 reference frame
q_rot_cw = euler_to_quaternion(0.0, 0.0, ang_rot)
q_rot_ccw = euler_to_quaternion(0.0, 0.0, -ang_rot) 
MOVEMENTS = {
    "home": [0.8, 0.0, 0.4, 0.0, 0.7071068, 0.0, 0.7071068], # Reset default pose
    "right": [0.0, -lat_displ, 0.0, 0.0, 0.0, 0.0, 0.0],  # Move right
    "left": [0.0, lat_displ, 0.0, 0.0, 0.0, 0.0, 0.0],  # Move left
    "down": [0.0, 0.0, -lat_displ, 0.0, 0.0, 0.0, 0.0],  # Move downwards
    "up": [0.0, 0.0, lat_displ, 0.0, 0.0, 0.0, 0.0],  # Move upwards
    "clockwise": [0.0, 0.0, 0.0, q_rot_cw[0], q_rot_cw[1], q_rot_cw[2], q_rot_cw[3]],  # Rotate clockwise around x axis
    "cclockwise": [0.0, 0.0, 0.0, q_rot_ccw[0], q_rot_ccw[1], q_rot_ccw[2], q_rot_ccw[3]],  # Rotate counter-clockwise 
    "vertical": "vertical", # Sequence downward upward
    "plane": "plane", # Sequence right left
    "fullcw": "fullcw", 
    "fullccw": "fullccw",
}

# ----- Global Variables -----
TARGET = Pose()
MOVEMENT = MOVEMENTS["home"]
EE_POSE = Pose()
keyboard_cmd = False
tf_buffer = None
is_home = False
is_rotation_cw = False
is_rotation_ccw = False
is_sequence = False
end_cw = False
end_ccw = False
frame_id_ee = "iiwa_link_7" 
delta_t = 0.0

# ----- Callbacks -----
def keyboard_callback(data):
    """Callback for keyboard commands

    Args:
        data (String): ROS `String` message
    """
    global MOVEMENT
    global keyboard_cmd
    global is_home
    global is_rotation_cw
    global is_rotation_ccw
    global is_sequence
    # Check if known target, otherwise keep default pose
    MOVEMENT = MOVEMENTS.get(data.data, MOVEMENTS["home"])
    if data.data == "home":
        is_home = True
    if (data.data == "clockwise") or (data.data == "fullcw"):
        is_rotation_cw = True
    if (data.data == "cclockwise") or (data.data == "fullccw"):
        is_rotation_ccw = True
    if (data.data == "vertical") or (data.data == "plane") or (data.data == "fullcw") or (data.data == "fullccw"):
        is_sequence = True
    keyboard_cmd = True

def ee_pose_callback(data):
    """Callback for end effector pose

    Args:
        data (Pose): ROS `Pose` message
    """
    global EE_POSE
    EE_POSE = data

# ----- Helper functions for pose and target computation-----
def to_pose():
    """Transform corresponding MOVEMENT in Pose object
    Return: 
        pose (Pose)
    """
    pose = Pose()
    pose.position.x = MOVEMENT[0]
    pose.position.y = MOVEMENT[1]
    pose.position.z = MOVEMENT[2]
    pose.orientation.x = MOVEMENT[3]
    pose.orientation.y = MOVEMENT[4]
    pose.orientation.z = MOVEMENT[5]
    pose.orientation.w = MOVEMENT[6]
    return pose

def sum_poses(pose1, pose2):
    """Sum the position and orientation fields between two poses

    Args:
        pose1 (Pose)
        pose2 (Pose)
    Return:
        sum_p (Pose)
    """
    sum_p = Pose()
    sum_p.position.x = pose1.position.x + pose2.position.x
    sum_p.position.y = pose1.position.y + pose2.position.y
    sum_p.position.z = pose1.position.z + pose2.position.z
    sum_p.orientation.x = pose1.orientation.x + pose2.orientation.x
    sum_p.orientation.y = pose1.orientation.y + pose2.orientation.y
    sum_p.orientation.z = pose1.orientation.z + pose2.orientation.z
    sum_p.orientation.w = pose1.orientation.w + pose2.orientation.w
    if is_rotation_cw or is_rotation_ccw or is_home:
        sum_p.orientation.x = pose2.orientation.x
        sum_p.orientation.y = pose2.orientation.y 
        sum_p.orientation.z = pose2.orientation.z 
        sum_p.orientation.w = pose2.orientation.w 
    
    return check_feasible(sum_p)

def to_world_frame(pose_mvt):
    """Use tf transform to express the desired rotation in the world frame

    Args:
        pose_mvt (Pose):pose associated to the desired movement in the ee frame
    Return:
        rel_target_pose (Pose): relative desired pose in the world frame
    """
    # Apply tf only if some rotation is involved 
    if (is_rotation_cw) or (is_rotation_ccw):
        try:
            pose_stamped = PoseStamped()
            pose_stamped.pose = pose_mvt
            pose_stamped.header.frame_id = frame_id_ee
            # Hack to ensure that the target transformation is always in future
            pose_stamped.header.stamp = rospy.Time.from_sec(rospy.Time.now().to_sec() + 0.5)
            output_pose_stamped = tf_buffer.transform(pose_stamped, base_link, timeout=rospy.Duration(1.0))
        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException):
            raise
        rel_target_pose = output_pose_stamped.pose
        rel_target_pose.position = pose_mvt.position # The desired position was already expressed in the world frame
    else:
        rel_target_pose = pose_mvt
   
    return  rel_target_pose

def check_feasible(p):
    """Check if boundary conditions are respected, constraint ee rotation to physical limits

    Args: 
        sum_p (Pose): computed sum of two poses
    Return:
        sum_p (Pose): saturate to limits if necessary 
    """
    global is_rotation_cw
    global is_rotation_ccw
    global end_cw
    global end_ccw
    
    if (is_rotation_cw):
        if (not end_ccw and np.abs(p.orientation.w + p.orientation.y) < 0.3) or end_cw:
            # Saturate all orientations to prevent overflow and keep quaternion normalized
            p.orientation.x = 0.7071068
            p.orientation.y = 0.0
            p.orientation.z = 0.7071068
            p.orientation.w = 0.0
            rospy.loginfo("cw saturation")
            end_cw = True
        end_ccw = False
        is_rotation_cw = False
    if (is_rotation_ccw):
        if (not end_cw and np.abs(p.orientation.w + p.orientation.y) < 0.3) or end_ccw:
            p.orientation.x = 0.7071068
            p.orientation.y = 0.0
            p.orientation.z = 0.7071068
            p.orientation.w = 0.0
            rospy.loginfo("ccw saturation")
            end_ccw = True
        end_cw = False
        is_rotation_ccw = False
    return p 

def get_error(p1, p2):
    """Return error between two poses

    Args:
        p1 (Pose): First Pose
        p2 (Pose): Second Pose

    Returns:
        Tuple[float, float]: Tuple of position and orientation error
            Position error: L2 distance between the poses
            Orientation error: Angle between the poses [rad]
    """
    pos_err = np.array([
        np.linalg.norm([p1.position.x - p2.position.x]),
        np.linalg.norm([p1.position.y - p2.position.y]),
        np.linalg.norm([p1.position.z - p2.position.z]),
    ])
    orient_err = np.abs(angle_between_quaternion(p1.orientation, p2.orientation))

    return pos_err, orient_err


# ----- Trajectories definition -----
def pub_target(rel_mvt):
    global TARGET
    TARGET = sum_poses(EE_POSE, rel_mvt)
    TARGET.orientation = normalize_quat(TARGET.orientation)
    publisher.publish(TARGET)

def reach_lin_target(nb_axis, radius_norm = 0.01):
    while not rospy.is_shutdown():
        err = get_error(EE_POSE, TARGET)
        
        if(np.linalg.norm(err[0][nb_axis]) < radius_norm): #define norm to change waypoint
            break

def reach_ang_target(angular_norm = 0.3):
    while not rospy.is_shutdown():
            err = get_error(EE_POSE, TARGET)
            if(np.linalg.norm(err[1]) < angular_norm):
                break

def vertical_mvt(rep, lat_displ, dt):
    global MOVEMENT
    mvt = Pose()
    mvt.position.x = 0
    mvt.position.y = 0
    mvt.position.z = 0
    mvt.orientation.x = EE_POSE.orientation.x
    mvt.orientation.y = EE_POSE.orientation.y
    mvt.orientation.z = EE_POSE.orientation.z
    mvt.orientation.w = EE_POSE.orientation.w
    for i in range(rep):
        if v_des != 0:
           
            mvt.position.z = -v_des
            TARGET = mvt
            publisher_v.publish(TARGET)
            t_init = time.time()
            while (time.time() - t_init < dt):
                pass
        
            mvt.position.z = v_des
           
            TARGET = mvt
            publisher_v.publish(TARGET)
            t_init = time.time()
            while (time.time() - t_init < dt + 0.03):
                pass
                
        else:
            
            rel_mvt = Pose()
            rel_mvt.position.z = -lat_displ
            pub_target(rel_mvt)
            reach_lin_target(2)
            rel_mvt.position.z = lat_displ
            pub_target(rel_mvt)
            reach_lin_target(2)
            
    #Reset position
    MOVEMENT = MOVEMENTS["home"]
    TARGET = to_pose()
    publisher.publish(TARGET)
  
        
def horizontal_mvt(rep, lat_displ, dt):
    global MOVEMENT
    mvt = Pose()
    mvt.position.x = 0
    mvt.position.y = 0
    mvt.position.z = 0
    mvt.orientation.x = EE_POSE.orientation.x
    mvt.orientation.y = EE_POSE.orientation.y
    mvt.orientation.z = EE_POSE.orientation.z
    mvt.orientation.w = EE_POSE.orientation.w
    for i in range(rep):
        if v_des != 0:
            
            mvt.position.y = -v_des
        
            TARGET = mvt
            publisher_v.publish(TARGET)
            t_init = time.time()
            while (time.time() - t_init < dt):
                pass

            mvt.position.y = v_des
            
            TARGET = mvt
            publisher_v.publish(TARGET)
            t_init = time.time()
            while (time.time() - t_init < dt):
                pass
            
        else:
            rel_mvt = Pose()
            rel_mvt.position.y = -lat_displ
            pub_target(rel_mvt)
            reach_lin_target(1)
            rel_mvt.position.y = lat_displ
            pub_target(rel_mvt)
            reach_lin_target(1)

    #Reset position
    MOVEMENT = MOVEMENTS["home"]
    TARGET = to_pose()
    publisher.publish(TARGET)
        

def full_cw(ang_rot, rep = 2):

    global MOVEMENT
    global is_rotation_cw
    global is_rotation_ccw

    for i in range(rep):
        is_rotation_cw = True #otherwise not computing tf
        rel_mvt = Pose()
        q_rot_cw = euler_to_quaternion(0.0, 0.0, ang_rot) 
        MOVEMENT = [0.0, 0.0, 0.0, q_rot_cw[0], q_rot_cw[1], q_rot_cw[2], q_rot_cw[3]]
        pose_mvt = to_pose()
        rel_mvt = to_world_frame(pose_mvt)
        pub_target(rel_mvt)
        reach_ang_target()

        is_rotation_ccw =True #otherwise not computing tf
        q_rot_ccw = euler_to_quaternion(0.0, 0.0, -ang_rot) 
        MOVEMENT = [0.0, 0.0, 0.0, q_rot_ccw[0], q_rot_ccw[1], q_rot_ccw[2], q_rot_ccw[3]]
        pose_mvt = to_pose()
        rel_mvt = to_world_frame(pose_mvt)
        pub_target(rel_mvt)
        reach_ang_target()
        
    #Reset position
    MOVEMENT = MOVEMENTS["home"]
    TARGET = to_pose()
    publisher.publish(TARGET)
  
        
def full_ccw(ang_rot, rep = 2):
    global MOVEMENT
    global is_rotation_cw
    global is_rotation_ccw

    for i in range(rep):
        is_rotation_ccw = True
        rel_mvt = Pose()
        q_rot_ccw = euler_to_quaternion(0.0, 0.0, -ang_rot) 
        MOVEMENT = [0.0, 0.0, 0.0, q_rot_ccw[0], q_rot_ccw[1], q_rot_ccw[2], q_rot_ccw[3]]
        pose_mvt = to_pose()
        rel_mvt = to_world_frame(pose_mvt)
        pub_target(rel_mvt)
        reach_ang_target()
       
        is_rotation_cw = True #otherwise not computing tf
        q_rot_cw = euler_to_quaternion(0.0, 0.0, ang_rot) 
        MOVEMENT = [0.0, 0.0, 0.0, q_rot_cw[0], q_rot_cw[1], q_rot_cw[2], q_rot_cw[3]]
        pose_mvt = to_pose()
        rel_mvt = to_world_frame(pose_mvt)
        pub_target(rel_mvt)
        reach_ang_target()

    #Reset position
    MOVEMENT = MOVEMENTS["home"]
    TARGET = to_pose()
    publisher.publish(TARGET)

def run_sequence():
    
    global delta_t
    if v_des != 0:
        delta_t = 0.3
    if MOVEMENT == "vertical":
        vertical_mvt(rep, lat_displ, delta_t)
    if MOVEMENT == "plane":
        horizontal_mvt(rep, lat_displ, delta_t)
    if MOVEMENT == "fullcw":
        full_cw(ang_rot)
    if MOVEMENT == "fullccw":
        full_ccw(ang_rot)

# ----- Main Script -----
if __name__ == '__main__':
    rospy.init_node('gazebo_iiwa_keyboard_cmd')
    rate = rospy.Rate(30)
    base_link = rospy.get_param("~base_link", "world")

    # Subscribers
    tf_buffer = tf2_ros.Buffer()
    listener = tf2_ros.TransformListener(tf_buffer)
    rospy.Subscriber("iiwa/ee_info/Pose", Pose, ee_pose_callback)
    rospy.Subscriber("iiwa/lib_cmd", String, keyboard_callback, queue_size=1)

    # Publishers
    publisher = rospy.Publisher("passive_control/pos_quat", Pose, queue_size=1)

    #Publisher to ee_vel_d
    publisher_v = rospy.Publisher("passive_control/vel_quat", Pose, queue_size=1)

    # Main loop
    while not rospy.is_shutdown():
        rep = 2
        if args.velocity and keyboard_cmd:
            if is_home:
                TARGET = to_pose()
                is_home = False
                TARGET.orientation = normalize_quat(TARGET.orientation)
                publisher.publish(TARGET)
            # Rotation
            elif is_rotation_cw or is_rotation_ccw:
                ang_rot = ang_rot/args.nb_sub_rot
                run_sequence()  
            
            else: #Linear movement
                v_des = args.v
                if v_des == 0:
                    rospy.loginfo('ATTENTION DESIRED VELOCITY IS 0')
                run_sequence()  
                
            is_sequence = False
            keyboard_cmd = False

        else:

            if(is_sequence): 
                run_sequence()  
                is_sequence = False
                keyboard_cmd = False
            elif(keyboard_cmd):
                if is_home:
                    TARGET = to_pose()
                    is_home = False
                else:
                    pose_mvt = to_pose()
                    rel_target_pose = to_world_frame(pose_mvt)
                    TARGET = sum_poses(EE_POSE, rel_target_pose)
                TARGET.orientation = normalize_quat(TARGET.orientation)
                publisher.publish(TARGET)
                keyboard_cmd = False
                rospy.loginfo(TARGET) # print for sanity check
           
        rate.sleep()
#!/usr/bin/env python # # Copyright (C) 2021 Learning Algorithms and Systems Laboratory, EPFL, Switzerland # Authors: # 	<NAME> (<EMAIL>) # # Website: lasa.epfl.ch # # This file is part of iiwa_gazebo. # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU General Public License as published by # the Free Software Foundation, either version 3 of the License, or # (at your option) any later version. # # This program is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the # GNU General Public License for more details. # # You should have received a copy of the GNU General Public License # along with this program. If not, see <http://www.gnu.org/licenses/>. # # ----- Argument parser ----- # ----- Helper functions for quaternions----- Normalise the quaternion      Args:         q (Quaternion)     Return:         q (Quaternion): normalized quaternion Return angle between two quaternions      Args:         p1 (Quaternion): First Quaternion         p2 (Quaternion): Second Quaternion      Returns:         float: Angle between the quaternions [rad] # Inv of q2 # norm [ vector component of q1 * q2^(-1) ] = sin(\theta / 2) # ----- Movements (predefined) ----- # m # radians # Desired rotation in the iiwa_link_7 reference frame # Reset default pose # Move right # Move left # Move downwards # Move upwards # Rotate clockwise around x axis # Rotate counter-clockwise # Sequence downward upward # Sequence right left # ----- Global Variables ----- # ----- Callbacks ----- Callback for keyboard commands      Args:         data (String): ROS `String` message # Check if known target, otherwise keep default pose Callback for end effector pose      Args:         data (Pose): ROS `Pose` message # ----- Helper functions for pose and target computation----- Transform corresponding MOVEMENT in Pose object     Return:          pose (Pose) Sum the position and orientation fields between two poses      Args:         pose1 (Pose)         pose2 (Pose)     Return:         sum_p (Pose) Use tf transform to express the desired rotation in the world frame      Args:         pose_mvt (Pose):pose associated to the desired movement in the ee frame     Return:         rel_target_pose (Pose): relative desired pose in the world frame # Apply tf only if some rotation is involved # Hack to ensure that the target transformation is always in future # The desired position was already expressed in the world frame Check if boundary conditions are respected, constraint ee rotation to physical limits      Args:          sum_p (Pose): computed sum of two poses     Return:         sum_p (Pose): saturate to limits if necessary # Saturate all orientations to prevent overflow and keep quaternion normalized Return error between two poses      Args:         p1 (Pose): First Pose         p2 (Pose): Second Pose      Returns:         Tuple[float, float]: Tuple of position and orientation error             Position error: L2 distance between the poses             Orientation error: Angle between the poses [rad] # ----- Trajectories definition ----- #define norm to change waypoint #Reset position #Reset position #otherwise not computing tf #otherwise not computing tf #Reset position #otherwise not computing tf #Reset position # ----- Main Script ----- # Subscribers # Publishers #Publisher to ee_vel_d # Main loop # Rotation #Linear movement # print for sanity check