# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/07_peripheral.ipynb.

# %% auto 0
__all__ = ['print_pygame_info', 'get_memorymaze_action_from_joystick', 'get_pinpad_action_from_joystick',
           'get_calvin_action_from_joystick']

# %% ../nbs/07_peripheral.ipynb 3
import pygame
import numpy as np

# Print all pygame information about connected peripherals
def print_pygame_info():
    pygame.init()
    pygame.joystick.init()
    joystick_count = pygame.joystick.get_count()
    print("Number of joysticks: {}".format(joystick_count))
    for i in range(joystick_count):
        joystick = pygame.joystick.Joystick(i)
        joystick.init()
        print("Joystick {}".format(i))
        print("    Name: {}".format(joystick.get_name()))
        print("    Axes: {}".format(joystick.get_numaxes()))
        print("    Balls: {}".format(joystick.get_numballs()))
        print("    Buttons: {}".format(joystick.get_numbuttons()))
        print("    Hats: {}".format(joystick.get_numhats()))

print_pygame_info()

# %% ../nbs/07_peripheral.ipynb 5
def get_memorymaze_action_from_joystick(joysticks):
    # make sure pygame is initialized
    pygame.init()
    
    leftx = joysticks[0].get_axis(0)
    lefty = joysticks[0].get_axis(1)
    leftradius, lefttheta = pygame.math.Vector2(leftx, lefty).as_polar()
    # trim the deadzone and convert from -180, 180 to 0, 360. This does a weird thing where the y-axis is flipped, that's why the numbers below are weird.
    deadzone_radius = 0.2
    if leftradius < deadzone_radius: lefttheta = None
    else: lefttheta += 180 

    # dictionary with bounds corresponding to each action
    theta_to_memorymaze_action = {
        # 0: (None, None), # noop
        1: (115, 65), # forward
        3: (200, 160), # right - swap order for moving across 0
        2: (20, 340), # left
        5: (160, 115), # forward_right
        4: (65, 20), # forward_left
    }

    action_to_string_memorymaze = {
        0: "noop",
        1: "forward",
        2: "left",
        3: "right",
        4: "forward_left",
        5: "forward_right",
    }

    chosen_action = None # noop default
    if lefttheta is not None:
        for action, (upper, lower) in theta_to_memorymaze_action.items():
            if lefttheta >= lower and lefttheta <= upper:
                chosen_action = action
            elif action_to_string_memorymaze[action] == "left" and (lefttheta >= lower or lefttheta <= upper): 
                # special case around 0
                chosen_action = action

    if chosen_action is not None:
        if chosen_action == 0: print(f"noop, rad {leftradius:1.2f} theta {lefttheta}")
        else: print(f"{action_to_string_memorymaze[chosen_action]}, {lefttheta}")

    return chosen_action



# %% ../nbs/07_peripheral.ipynb 6
def get_pinpad_action_from_joystick(joysticks):
    # make sure pygame is initialized
    pygame.init()
    pygame.event.pump()
    
    leftx = joysticks[0].get_axis(0)
    lefty = joysticks[0].get_axis(1)
    leftradius, lefttheta = pygame.math.Vector2(leftx, lefty).as_polar()
    # trim the deadzone and convert from -180, 180 to 0, 360. This does a weird thing where the y-axis is flipped, that's why the numbers below are weird.
    deadzone_radius = 0.2
    if leftradius < deadzone_radius: lefttheta = None
    else: lefttheta += 180 

    # dictionary with bounds corresponding to each action
    theta_to_pinpad_action = {
        # 0: (None, None), # noop
        2: (135, 45), # forward
        3: (225, 135), # right - swap order for moving across 0
        4: (45, 315), # left
        1: (315, 225), # bottom
        # 4: (65, 20), # forward_left
    }

    action_to_string_pinpad = {
        0: "noop",
        1: "down",
        2: "fwd",
        3: "right",
        4: "left",
    }

    chosen_action = None # noop default
    if lefttheta is not None:
        for action, (upper, lower) in theta_to_pinpad_action.items():
            if lefttheta >= lower and lefttheta <= upper:
                chosen_action = action
            elif action_to_string_pinpad[action] == "left" and (lefttheta >= lower or lefttheta <= upper): 
                # special case around 0
                chosen_action = action

    if chosen_action is not None:
        if chosen_action == 0: print(f"noop, rad {leftradius:1.2f} theta {lefttheta}")
        else: print(f"{action_to_string_pinpad[chosen_action]}, {lefttheta}")

    return chosen_action


# t0 = time.time()

# pygame.init()

# # Set up the drawing window (screen)
# screen = pygame.display.set_mode([500, 500])

# # initialize xbox controller
# joysticks = [pygame.joystick.Joystick(x) for x in range(pygame.joystick.get_count())]

# while time.time() - t0 < 10:
#     get_pinpad_action_from_joystick(joysticks)


# %% ../nbs/07_peripheral.ipynb 7
def get_calvin_action_from_joystick(joysticks):

    def get_joystick_xytheta(joystick, axis=0):
        '''
        Returns the x, y, radius, and theta of the joystick.
        default left joystick
        '''
        x = joysticks[0].get_axis(axis)
        y = joysticks[0].get_axis(axis+1)
        r, theta = pygame.math.Vector2(x, y).as_polar()
        # trim the deadzone and convert from -180, 180 to 0, 360. This does a weird thing where the y-axis is flipped, that's why the numbers below are weird.
        deadzone_radius = 0.2
        if r < deadzone_radius: theta = None
        else: theta += 180
        return x, y, r, theta
    # make sure pygame is initialized
    pygame.init()

    # get the joystick values
    leftx, lefty, leftradius, lefttheta = get_joystick_xytheta(None, axis=0)
    rightx, righty, rightradius, righttheta = get_joystick_xytheta(None, axis=3)


    dx = float(leftx) if lefttheta is not None else 0.
    dy = float(lefty) if lefttheta is not None else 0.
    dz = float(righty) if righttheta is not None else 0.
    # print(f"dx: {dx:1.2f}, dy: {dy:1.2f}, dz: {dz:1.2f}")
    dpitch = float(rightx) if righttheta is not None else 0.

    trigger0 = joysticks[0].get_axis(2)
    trigger1 = joysticks[0].get_axis(5)
    droll = float(trigger0) if trigger0 > 0.1 else 0.
    dyaw = float(trigger1) if trigger1 > 0.1 else 0.

    ret = (dx, dy, dz, droll, dpitch, dyaw, 1) # 0 is the gripper
    print(ret)
    return ret

    theta_to_memorymaze_action = {
        # 0: (None, None), # noop
        1: (115, 65), # forward
        3: (200, 160), # right - swap order for moving across 0
        2: (20, 340), # left
        5: (160, 115), # forward_right
        4: (65, 20), # forward_left
    }

    action_to_string_memorymaze = {
        0: "noop",
        1: "forward",
        2: "left",
        3: "right",
        4: "forward_left",
        5: "forward_right",
    }


    return leftx, lefty, rightx, righty

    # dictionary with bounds corresponding to each action
    theta_to_memorymaze_action = {
        # 0: (None, None), # noop
        1: (115, 65), # forward
        3: (200, 160), # right - swap order for moving across 0
        2: (20, 340), # left
        5: (160, 115), # forward_right
        4: (65, 20), # forward_left
    }

    action_to_string_memorymaze = {
        0: "noop",
        1: "forward",
        2: "left",
        3: "right",
        4: "forward_left",
        5: "forward_right",
    }

    chosen_action = None # noop default
    if lefttheta is not None:
        for action, (upper, lower) in theta_to_memorymaze_action.items():
            if lefttheta >= lower and lefttheta <= upper:
                chosen_action = action
            elif action_to_string_memorymaze[action] == "left" and (lefttheta >= lower or lefttheta <= upper): 
                # special case around 0
                chosen_action = action

    if chosen_action is not None:
        if chosen_action == 0: print(f"noop, rad {leftradius:1.2f} theta {lefttheta}")
        else: print(f"{action_to_string_memorymaze[chosen_action]}, {lefttheta}")

    return chosen_action


