import umath
from pybricks.pupdevices import Motor, ColorSensor, UltrasonicSensor, ForceSensor
from pybricks.parameters import Button, Color, Direction, Port, Side, Stop
from pybricks.tools import run_task, multitask
from pybricks.tools import wait, StopWatch
from pybricks.robotics import DriveBase
from pybricks.hubs import PrimeHub

# Initialize hub and devices
hub = PrimeHub()
left_motor = Motor(Port.A, Direction.COUNTERCLOCKWISE)
right_motor = Motor(Port.B,Direction.CLOCKWISE) # Specify default direction
left_arm = Motor(Port.C, Direction.CLOCKWISE, [[12,36]],[[12,20,24]] ) # Specify default direction
right_arm = Motor(Port.D, Direction.CLOCKWISE,[[12,36]],[[12,20,24]]) #Added gear train list for gear ration
lazer_ranger = UltrasonicSensor(Port.E)
color_sensor = ColorSensor(Port.F)

# DriveBase configuration
WHEEL_DIAMETER = 62.4  # mm (adjust for your wheels)
AXLE_TRACK = 150  # mm (distance between wheels)
drive_base = DriveBase(left_motor, right_motor, WHEEL_DIAMETER, AXLE_TRACK)
drive_base.settings(600, 500, 300, 200)
drive_base.use_gyro(True)

def set_default_speed():
    drive_base.settings(600, 500, 300, 200)

def set_speed(straight_speed, st_acc, turn_speed, turn_acc):
    drive_base.settings(straight_speed, st_acc, turn_speed, turn_acc)

async def main():

    await drive_base.straight(110)
    await drive_base.turn(-45)
    await drive_base.straight(240)
    await right_arm.run_angle(300,120)
    await drive_base.straight(290)
    await drive_base.turn(50)
    await drive_base.straight(300)
    await drive_base.turn(40)
    set_speed(100,200,300,200)
    await drive_base.straight(130)
    await drive_base.straight(-115)
    await drive_base.turn(87)
    await drive_base.straight(230)
    set_speed(600, 500, 300, 200)
    await drive_base.straight(-135)
    await multitask(
        drive_base.turn(-40),
        right_arm.run_angle(100,-110)
    )
    set_speed(100,200,300,200)
    await drive_base.straight(220)
    await right_arm.run_angle(300,85)
    await wait(500)
    await drive_base.straight(142)
    await drive_base.arc(110,None,-166)
    await drive_base.straight(140)
    await drive_base.straight(-150)
    set_speed(600,500,300,200)
    await drive_base.straight(150)
    await drive_base.arc(800,None,-900)

run_task(main())