from time import sleep
import robot

POWER = 70

TURN_T = 0.068 # 10 degrees
DRIVE_T = 0.22 # 10 centimeter

RIGHT_WHEEL_OFFSET = 5

CLOCKWISE_OFFSET = 1

def wiggle(arlo):
    arlo.go_diff(50, 53, 1, 0)
    sleep(TURN_T * CLOCKWISE_OFFSET * 2)
    arlo.stop()
    sleep(0.3)
    arlo.go_diff(50, 53, 0, 1)
    sleep(TURN_T * 4)
    arlo.stop()
    sleep(0.3)
    arlo.go_diff(50, 53, 1, 0)
    sleep(TURN_T * CLOCKWISE_OFFSET * 2)
    arlo.stop()

def clockwise_square(arlo):
    arlo.go_diff(POWER, POWER + RIGHT_WHEEL_OFFSET, 0, 1)
    sleep(TURN_T * 0.5)
    arlo.stop()
    sleep(0.5)

    for _ in range(3):
        for _ in range(4):
            # Driving forward
            arlo.go_diff(POWER, POWER + RIGHT_WHEEL_OFFSET, 1, 1)
            sleep(DRIVE_T * 10)
            arlo.stop()

            sleep(1)

            # Turning 90 degrees
            arlo.go_diff(POWER, POWER, 1, 0)
            sleep(TURN_T * 10 * CLOCKWISE_OFFSET)
            arlo.stop()
            sleep(0.2)
            arlo.go_diff(POWER, POWER + RIGHT_WHEEL_OFFSET, 0, 1)
            sleep(TURN_T)
            arlo.stop()

            sleep(1)
        sleep(1.5)

    for _ in range(4):
        wiggle(arlo)

def counter_square(arlo):
    arlo.go_diff(POWER, POWER + RIGHT_WHEEL_OFFSET, 0, 1)
    sleep(TURN_T * 0.5)
    arlo.stop()
    sleep(0.5)

    for _ in range(12):
        # Driving forward
        arlo.go_diff(POWER, POWER + RIGHT_WHEEL_OFFSET, 1, 1)
        sleep(DRIVE_T * 10)
        arlo.stop()

        sleep(0.5)

        # Turning 90 degrees
        arlo.go_diff(POWER, POWER + RIGHT_WHEEL_OFFSET, 0, 1)
        sleep(TURN_T * 10)
        arlo.stop()
        sleep(0.2)
        arlo.go_diff(POWER, POWER + RIGHT_WHEEL_OFFSET, 1, 0)
        sleep(TURN_T)
        arlo.stop()

        sleep(0.5)

def main():
    # Initializes the robot and runs the
    arlo = robot.Robot()

    direction = int(input("clockwise (1)/counter (2):"))

    try:
        if direction == 1:
            clockwise_square(arlo)
        else:
            counter_square(arlo)

    except KeyboardInterrupt:
        arlo.stop()


if __name__ == "__main__":

    main()