import rclpy
from rclpy.node import Node

from std_srvs.srv import Trigger

from std_msgs.msg import Float32
from sensor_msgs.msg import Joy
from geometry_msgs.msg import TwistStamped
from control_msgs.msg import JointJog
from moveit_msgs.msg import PlanningScene, PlanningSceneWorld

from enum import Enum


# Mappings for a DualShock4 (PS4) controller
class Buttons(Enum):
    CROSS = 0
    CIRCLE = 1
    TRIANGLE = 2
    SQUARE = 3
    L1 = 4
    R1 = 5
    L2 = 6
    R2 = 7
    SELECT = 8
    START = 9
    CENTER = 10
    L3 = 11
    R3 = 12

class Axes(Enum):
    LEFT_X = 0
    LEFT_Y = 1
    L1 = 2
    RIGHT_X = 3
    RIGHT_Y = 4
    R1 = 5
    DPAD_UPDOWN = 6
    DPAD_LEFTRIGHT = 7

class JoyMove(Node):
    def __init__(self):
        super().__init__('joystickmove')

        # Get the logger
        self._logger = self.get_logger()
        self._logger.info('joystick init')

        # Set up topics for publishing
        JOY_TOPIC = '/joy'
        TWIST_TOPIC = '/servo_server/delta_twist_cmds'
        JOINT_TOPIC = '/servo_server/delta_joint_cmds'

        # Define some constants
        EE = 'joint_6'
        BASE = 'base_link'

        # Create the needed subscribers and publishers for moveit_servo
        self.create_subscription(Joy, '/joy', self._joy_cb, rclpy.qos.qos_profile_sensor_data)

        self.twist_pub = self.create_publisher(TwistStamped, TWIST_TOPIC, 10);
        self.joint_pub = self.create_publisher(JointJog, JOINT_TOPIC, 10);
        self.collision_pub = self.create_publisher(PlanningScene, "/planning_scene", 10);

        # Start a moveit_servo client
        self.servo_client = self.create_client(Trigger, '/servo_server/start_servo')
        # Check if the a service is available
        while not self.servo_client.wait_for_service(timeout_sec=1.0):
            self.get_logger().info('servo_server not available, waiting again...')
        self._logger.info('servo_server connected.')

        # Create a request to tell the server the joystick exists
        req = Trigger.Request()
        self.servo_client.call_async(req)

    # Callback for any button or stick that is pressed on the joystick
    def _joy_cb(self, msg):
        axes = msg.axes
        buttons = msg.buttons

        # If you press main buttons or dpad arrows
        if buttons[Buttons.CROSS.value] or buttons[Buttons.CIRCLE.value] or buttons[Buttons.TRIANGLE.value] \
            or buttons[Buttons.SQUARE.value] or axes[Axes.DPAD_UPDOWN.value] or axes[Axes.DPAD_LEFTRIGHT.value]:

            # combine default values and whatever you have pushed to create a JointJog object
            joint = JointJog()

            joint.joint_names = ['joint_1', 'joint_2',
                                 'joint_6', 'joint_5']
            joint.velocities = [float(axes[Axes.DPAD_UPDOWN.value]), float(axes[Axes.DPAD_LEFTRIGHT.value]),
                                float(buttons[Buttons.CIRCLE.value] - buttons[Buttons.CROSS.value]), float(buttons[Buttons.SQUARE.value] - buttons[Buttons.TRIANGLE.value])]
            joint.duration = 0.1

            # Add a timestamp so we know when the action actually took place
            joint.header.stamp = self.get_clock().now().to_msg()
            joint.header.frame_id = "link_6"

            # Publish the JointJog to servo_server
            self.joint_pub.publish(joint)
            return

        # If you pressed a trigger or stick
        ts = TwistStamped()

        # Define linear depth from right stick
        ts.twist.linear.z = axes[Axes.RIGHT_Y.value]
        ts.twist.linear.y = axes[Axes.RIGHT_X.value]

        # Define linear X from trigger analog values
        lin_x_right = -0.5 * axes[Axes.R1.value] # - AXIS_DEFAULTS.at(R1));
        lin_x_left = 0.5 * axes[Axes.L1.value] # - AXIS_DEFAULTS.at(L1));
        ts.twist.linear.x = lin_x_right + lin_x_left

        # Define angular up and down from left stick
        ts.twist.angular.y = axes[Axes.LEFT_Y.value]
        ts.twist.angular.x = axes[Axes.LEFT_X.value]

        # Left and Right bumpers to control roll
        roll_positive = buttons[Buttons.R2.value]
        roll_negative = -1 * buttons[Buttons.L2.value]
        ts.twist.angular.z = float(roll_positive + roll_negative)

        # Add a timestamp so we know when the action actually took place
        ts.header.stamp = self.get_clock().now().to_msg()
        ts.header.frame_id = "base_link"

        # Publish the Twist to servo_server
        self.twist_pub.publish(ts)


def main(args=None):
    rclpy.init(args=args)

    # Start the node
    joy_move = JoyMove()
    rclpy.spin(joy_move)

    # Kill the node
    joy_move.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()
