xbox_controller.py

from math import *

import cozmo
from cozmo.util import distance_mm

from xinput import *

directional_pad_speeds = {
    # up, down, left, right
    GAMEPAD_DPAD_UP: (100, 100),
    GAMEPAD_DPAD_DOWN: (-100, -100),
    GAMEPAD_DPAD_LEFT: (-100, 100),
    GAMEPAD_DPAD_RIGHT: (100, -100),
    # diagonals
    GAMEPAD_DPAD_UP | GAMEPAD_DPAD_LEFT: (50, 100),
    GAMEPAD_DPAD_UP | GAMEPAD_DPAD_RIGHT: (100, 50),
    GAMEPAD_DPAD_DOWN | GAMEPAD_DPAD_LEFT: (-50, -100),
    GAMEPAD_DPAD_DOWN | GAMEPAD_DPAD_RIGHT: (-100, -50),
}


def normalize_stick(x, y):
    """
    Normalize input values for left and right sticks
    :param x: x value generated by the controller
    :param y: y value generated by the controller
    :return: a tuple containing (normalized x, normalized_y, magnitude)
    """
    # print("x = {0}, y = {1}".format(x, y))
    # determine how far the controller is pushed
    magnitude = sqrt(x * x + y * y) + 0.01  # 0.01 to avoid division by zero when x = y = 0
    # determine the direction the controller is pushed
    normalized_x, normalized_y = x / magnitude, y / magnitude
    # check if the controller is outside a circular dead zone
    if magnitude > GAMEPAD_LEFT_THUMB_DEADZONE:
        magnitude = min([magnitude, GAMEPAD_THUMB_MAX])
        magnitude -= GAMEPAD_LEFT_THUMB_DEADZONE
    else:
        normalized_x, normalized_y, magnitude = 0.0, 0.0, 0.0

    return normalized_x, normalized_y, magnitude


def check_controller_state(robot: cozmo.robot.Robot, state):
    # face buttons
    # lift

    if state['buttons'] == GAMEPAD_B:
        robot.move_lift(1.0)
    elif state['buttons'] == GAMEPAD_A:
        robot.move_lift(-1.0)
    else:
        robot.move_lift(0)
    # head
    if state['buttons'] == GAMEPAD_Y:
        cube = robot.world.wait_for_observed_light_cube(timeout=30)
        if cube:
            robot.go_to_object(cube, distance_mm(200.0)).wait_for_completed()

            robot.pickup_object(cube, num_retries=0).wait_for_completed()

    # left stick
    left_x, left_y, left_magnitude = normalize_stick(state['l_thumb_x'], state['l_thumb_y'])
    # print("left :{0}, {1}, {2}".format(left_x, left_y, left_magnitude))

    # directional pad buttons

    if state['left_trigger'] > 0 or state['right_trigger'] > 0:
        robot.drive_wheels(state['left_trigger'], state['right_trigger'])
    else:
        (left_speed, right_speed) = directional_pad_speeds.get(state['buttons'] & 0xFF, (0, 0))
        if left_speed == 0.0 and right_speed == 0.0:
            if left_magnitude != 0.0:
                # up
                if left_y >= 0.75 and abs(left_x) < 0.25:
                    robot.drive_wheels(100, 100)
                # down
                elif left_y <= -0.75 and abs(left_x) < 0.25:
                    robot.drive_wheels(-100, -100)
                # left
                elif abs(left_y) < 0.25 and left_x <= -0.75:
                    robot.drive_wheels(-100, 100)
                # right
                elif abs(left_y) < 0.25 and left_x >= 0.75:
                    robot.drive_wheels(100, -100)
                # up + left
                elif left_y >= 0.50 and left_x <= -0.50:
                    robot.drive_wheels(50, 100)
                # up + right
                elif left_y >= 0.50 and left_x >= 0.50:
                    robot.drive_wheels(100, 50)
                # down + left
                elif left_y <= -0.50 and left_x <= -0.50:
                    robot.drive_wheels(-50, -100)
                # down + right
                elif left_y <= -0.50 and left_x >= 0.50:
                    robot.drive_wheels(-100, -50)
                else:
                    pass
            else:
                robot.drive_wheels(0, 0)
        else:
            robot.drive_wheels(left_speed, right_speed)


def cozmo_program(robot: cozmo.robot.Robot):
    joysticks = XInputJoystick.enumerate_devices()

    if joysticks:
        print("Number of connected controllers: {0}".format(len(joysticks)))
    else:
        print("No controller is connected. Please connect xbox controller.")
        sys.exit(0)
    # use only the first controller
    joystick = joysticks[0]

    while True:
        state = joystick.get_state()
        # print(state)
        check_controller_state(robot, state)
        time.sleep(.01)


cozmo.run_program(cozmo_program)