player.py

import socket
from math import *
from typing import List

import cozmo
from cozmo.util import distance_mm

from common.message_forwarder import start_connection, receive_message
from common.setup import get_team_colors
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(0.3)
    elif state['buttons'] == GAMEPAD_A:
        robot.move_lift(-0.3)
    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):
    """
    Main entry for running the player logic. This runs both the xbox controller
    functionality and checks for the exit message over the network when a team wins.

    :param robot: player robot in the game
    """
    # get number of teams playing in the game
    while True:
        try:
            teams: int = int(input("How many teams are playing?"))
        except ValueError:
            print("Invalid input type")
            continue
        if teams < 2:
            print("Must be between 2 and 3")
            continue
        elif teams > 3:
            print("Must be between 1 and 3")
            continue
        else:
            break

    # get the id of the team the judge is on
    while True:
        try:
            team_id: int = int(input("Which team is this player on?"))
        except ValueError:
            print("Invalid input type")
            continue
        if team_id < 1:
            print("Must be between 1 and 3")
            continue
        elif team_id > 3:
            print("Must be between 1 and 3")
            continue
        else:
            break

    # get the corresponding team colors and opponent colors
    team_colors, opponent_colors = get_team_colors(teams)

    # set backpack color
    robot.set_all_backpack_lights(team_colors[team_id])

    # set head angle so Cozmo can identify cubes
    robot.set_head_angle(cozmo.util.Angle(degrees=0))
    # establish connection to the network and message retrieval
    connection: socket.socket = start_connection("10.0.1.10", 5000)
    message: List[str] = []

    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 'Exit' not in message:
        state = joystick.get_state()
        check_controller_state(robot, state)
        message = receive_message(connection)

    # play the appropriate robot emotion based on who won the game
    if int(message[1]) == team_id:
        robot.play_anim_trigger(cozmo.anim.Triggers.CodeLabCelebrate).wait_for_completed()
    else:
        robot.play_anim_trigger(cozmo.anim.Triggers.CodeLabUnhappy).wait_for_completed()


if __name__ == '__main__':
    cozmo.run_program(cozmo_program)