import time
from machine import UART

# ARn
ARNI_INITIAL_VELOCITY = 0.5

# ODrive Constants
AXIS_STATE_UNDEFINED = 0
AXIS_STATE_IDLE = 1
AXIS_STATE_STARTUP_SEQUENCE = 2
AXIS_STATE_FULL_CALIBRATION_SEQUENCE = 3
AXIS_STATE_MOTOR_CALIBRATION = 4
AXIS_STATE_SENSORLESS_CONTROL = 5
AXIS_STATE_ENCODER_INDEX_SEARCH = 6
AXIS_STATE_ENCODER_OFFSET_CALIBRATION = 7
AXIS_STATE_CLOSED_LOOP_CONTROL = 8

# DABBLE BITS
# BUTTONS
DABBLE_START_BIT = 0
DABBLE_SELECT_BIT = 1
DABBLE_TRIANGLE_BIT = 2
DABBLE_CIRCLE_BIT = 3
DABBLE_CROSS_BIT = 4
DABBLE_SQUARE_BIT = 5
# ARROW PAD
DABBLE_UP_BIT = 0
DABBLE_DOWN_BIT = 1
DABBLE_LEFT_BIT = 2
DABBLE_RIGHT_BIT = 3

# Globals
uart = UART(2, baudrate=115200, tx=17, rx=16)
current_global_velocity = ARNI_INITIAL_VELOCITY  # Will be changed by controller during operation

class RobotController:

    def __init__(self, robot):
        self.robot = robot

    def interpret(self, dabble_bytes):
        """
        Interprets a tuples of bytes received from the Dabble BLE gamepad module.

        Translates data into ArniRobot commands.

        :param values: A two byte tuple
            values[0] contains button data
            values[1] contains arrow pad data
        :return:
        """
        self._interpret_buttons(dabble_bytes[0])
        self._interpret_arrows(dabble_bytes[1])

    def _interpret_buttons(self, button_byte):
        if button_byte & (1 << DABBLE_START_BIT):
            self.robot.init_motors()
        elif button_byte & (1 << DABBLE_SELECT_BIT):
            if self.robot.active:
                self.robot.set_idle()
            else:
                self.robot.set_active()

    def _interpret_arrows(self, arrow_byte):
        if arrow_byte & (1 << DABBLE_UP_BIT):
            self.robot.forward()
        elif arrow_byte & (1 << DABBLE_DOWN_BIT):
            self.robot.reverse()
        elif arrow_byte & (1 << DABBLE_LEFT_BIT):
            self.robot.left()
        elif arrow_byte & (1 << DABBLE_RIGHT_BIT):
            self.robot.right()
        else:
            self.robot.stop()

class ArniRobot:

    def __init__(self, uart, vel=None):
        self._uart = uart
        self._vel = vel if vel else ARNI_INITIAL_VELOCITY
        self.active = False
        self.initialized = False


    def get(self, value):
        """Read value from odrive config.

        Example:
            arni.get('vbus_voltage')
        """
        self._uart.write("r {0}\n".format(value))
        if not self._uart.any():
            time.sleep_ms(100)
        value = self._uart.read()
        if value:
            value = value.decode('ascii')
        return value or ''

    def init_motors(self):
        """Runs index ops and sets control loop.

        Must be run with the motors not under load.
        """
        # AXIS_STATE_ENCODER_OFFSET_CALIBRATION = 7
        self._uart.write('w axis0.requested_state 7\nw axis1.requested_state 7\n')
        self.initialized = True

    def forward(self):
        """Move forward."""
        return self._uart.write('v 0 {0} 0\nv 1 -{0} 0\n'.format(self._vel))

    def reverse(self):
        """Reverse move."""
        return self._uart.write('v 0 -{0} 0\nv 1 {0} 0\n'.format(self._vel))

    def stop(self):
        """All stop."""
        return self._uart.write('v 0 0 0\nv 1 0 0\n')

    def right(self):
        """Hard turn right."""
        return self._uart.write('v 0 0 0\nv 1 -{0} 0\n'.format(self._vel))

    def left(self):
        """Hard turn left."""
        return self._uart.write('v 0 {0} 0\nv 1 0 0\n'.format(self._vel))

    def right_forward(self):
        """Right turn while moving forward."""
        return self._uart.write('v 0 {0} 0\nv 1 -{1} 0\n'.format(self._vel / 2, self._vel))

    def right_reverse(self):
        """Right turn while moving in reverse."""
        return self._uart.write('v 0 -{0} 0\nv 1 {1} 0\n'.format(self._vel / 2, self._vel))

    def left_forward(self):
        """Left turn while moving forward."""
        return self._uart.write('v 0 {0} 0\nv 1 -{1} 0\n'.format(self._vel, self._vel / 2))

    def left_reverse(self):
        """Right turn while moving in reverse."""
        return self._uart.write('v 0 -{0} 0\nv 1 {0} 0\n'.format(self._vel, self._vel / 2))

    def set_idle(self):
        """Set motor control to idle."""
        self.stop()
        self.active = False
        return self._uart.write('w axis0.requested_state {0})\nw axis1.requested_state {0}\n'.format(AXIS_STATE_IDLE))

    def set_active(self):
        """Set motor control to closed loop control."""
        if not self.initialized:
            print("Robot motors have not been initialized. Initialize motors before entering control loop.")

        self.active = True
        return self._uart.write('w axis0.requested_state {0})\nw axis1.requested_state {0}\n'.format(AXIS_STATE_CLOSED_LOOP_CONTROL))

    def set_velocity(self, velocity):
        """Set motor velocity."""
        self._vel = velocity
        return self._vel

arni = ArniRobot(uart)
controller = RobotController(arni)