steering_wheel.py

""" Steering wheel """

import pygame

JOYSTICK_DIRECTION=0
JOYSTICK_THROTTLE=1
JOYSTICK_BRAKES=2
JOYSTICK_REVERSE=3

INPUT_OFFSET = 2.0

IDLE=0
FORWARD=1
REVERSE=2
FORWARD_LEFT=3
FORWARD_RIGHT=4
LEFT=5
RIGHT=6
BRAKE=7
REVERSE_LEFT=8
REVERSE_RIGHT=9
ACTION_SPACE=10

class ManualSteeringWheel:
    """ Steering wheel """
    axis_count = 0.0
    joystick = None
    direction = 0.0
    throttle = 0.0
    brakes = 0.0
    reverse = False


    def __init__(self):
        self.joystick = pygame.joystick.Joystick(0)
        self.joystick.init()
        self.axis_count = self.joystick.get_numaxes()

    def get_direction(self, update=True):
        """ Update and return direction. """
        if update:
            pygame.event.get()
        self.direction = self.joystick.get_axis(JOYSTICK_DIRECTION)
        return self.direction
    
    def get_throttle(self, update=True):
        """ Update and return throttle."""
        if update:
            pygame.event.get()
        self.throttle = round((INPUT_OFFSET - (self.joystick.get_axis(JOYSTICK_THROTTLE) + 1.0)) / INPUT_OFFSET, 3) 
        return self.throttle

    def get_brakes(self, update=True):
        """ Update and return brakes. """
        if update:
            pygame.event.get()
        self.brakes = (INPUT_OFFSET - (self.joystick.get_axis(JOYSTICK_BRAKES) + 1.0)) / INPUT_OFFSET
        self.brakes = round(self.brakes, 3)
        return self.brakes

    def update_controls(self):
        """ Poll for pygame events. """
        self.get_direction(update=False)
        self.get_throttle(update=False)
        self.get_brakes(update=False)

    def action_to_controls(action):
        if action == IDLE:
            return [0.0, 0.0, 0.0, False]

        if action == FORWARD:
            return [0.6, 0.0, 0.0, False]

        if action == REVERSE:
            return [0.5, 0.0, 0.0, True]

        if action == FORWARD_LEFT:
            return [0.4, -0.4, 0.0, False]

        if action == FORWARD_RIGHT:
            return [0.4, 0.4, 0.0, False]

        if action == LEFT:
            return [0.0, -0.4, 0.0, False]

        if action == RIGHT:
            return [0.0, 0.4, 0.0, False]

        if action == BRAKE:
            return [0.0, 0.0, 0.5, False]

        if action == REVERSE_LEFT:
            return [0.4, -0.4, 0.0, True]

        if action == REVERSE_RIGHT:
            return [0.4, 0.4, 0.0, True]
        
    def get_action(self):
        action = IDLE
        if self.throttle > 0.5 and self.brakes < 0.5 and abs(self.direction) < 0.1 and not self.reverse:
            action = FORWARD
        elif self.throttle > 0.5 and self.brakes < 0.5 and self.direction < -0.2 and not self.reverse:
            action = FORWARD_LEFT
        elif self.throttle > 0.5 and self.brakes < 0.5 and self.direction > 0.2 and not self.reverse:
            action = FORWARD_RIGHT
        elif self.throttle > 0.5 and self.brakes < 0.5 and abs(self.direction) < 0.2 and self.reverse:
            action = REVERSE
        elif self.throttle < 0.5 and self.brakes < 0.5 and self.direction < -0.2:
            action = LEFT
        elif self.throttle < 0.5 and self.brakes < 0.5 and self.direction > 0.2:
            action = RIGHT
        elif self.brakes > 0.5:
            action = BRAKE
        elif self.throttle > 0.5 and self.brakes < 0.5 and self.direction < -0.2 and self.reverse:
            action = REVERSE_LEFT
        elif self.throttle > 0.5 and self.brakes < 0.5 and self.direction > 0.2 and self.reverse:
            action = REVERSE_RIGHT
        return action


class Controller(ManualSteeringWheel):
    running = True

    def is_running(self):
        return self.running

    def on_update(self):
        for event in pygame.event.get():
            if event.type == pygame.KEYUP:
                if event.key == pygame.K_r:
                    self.reverse = True
                elif event.key == pygame.K_d:
                    self.reverse = False
                elif event.key == pygame.K_q:
                    self.running = False
                    
        self.update_controls()

if __name__ == '__main__':
    pygame.init()
    clock = pygame.time.Clock()
    sw = ManualSteeringWheel()

    while True:
        # Receive all occured events
        print("")
        print("Direction: " + str(sw.get_direction()))
        print("Speed:     " + str(sw.get_throttle()))
        print("Brakes:    " + str(sw.get_brakes()))
        clock.tick(2)