from machine import Pin, PWM, I2C, time_pulse_us
import ssd1306
import time

i2c = I2C(0, sda=Pin(4), scl=Pin(5), freq=400000)
oled = ssd1306.SSD1306_I2C(128, 64, i2c)

trig = Pin(21, Pin.OUT)
echo = Pin(22, Pin.IN)

SERVO_FREQ = 50

servo_left  = PWM(Pin(18), freq=SERVO_FREQ)
servo_right = PWM(Pin(19), freq=SERVO_FREQ)

def set_angle(servo, angle):
    duty = int(40 + (angle / 180) * (115 - 40))
    servo.duty(duty)

def measure_distance():
    trig.value(0)
    time.sleep_us(2)
    trig.value(1)
    time.sleep_us(10)
    trig.value(0)

    duration = time_pulse_us(echo, 1, 30000)

    if duration < 0:
        return -1

    return round(duration / 58.0, 1)

def sweep_forward():
    for angle in range(0, 181, 10):
        set_angle(servo_left, angle)
        set_angle(servo_right, 180 - angle)
        time.sleep_ms(10)

    for angle in range(180, -1, -10):
        set_angle(servo_left, angle)
        set_angle(servo_right, 180 - angle)
        time.sleep_ms(10)

def sweep_turn_right():
    for angle in range(0, 181, 10):
        set_angle(servo_left, angle)
        set_angle(servo_right, 90)
        time.sleep_ms(10)

    for angle in range(180, -1, -10):
        set_angle(servo_left, angle)
        set_angle(servo_right, 90)
        time.sleep_ms(10)

def move_robot(distance):
    if distance < 0:
        set_angle(servo_left, 90)
        set_angle(servo_right, 90)
        return "ERROR"

    if distance > 30:
        sweep_forward()
        return "FORWARD"

    elif distance >= 10:
        set_angle(servo_left, 0)
        set_angle(servo_right, 180)
        return "STOP"

    else:
        sweep_turn_right()
        return "TURN RIGHT"

def update_display(distance, state):
    oled.fill(0)
    oled.text("MOBILE ROBOT", 16, 0)
    oled.text("----------------", 0, 10)

    if distance < 0:
        oled.text("Dist: ERROR", 0, 22)
    else:
        oled.text("Dist: {} cm".format(distance), 0, 22)

    oled.text("State:", 0, 40)
    oled.text(state, 0, 52)
    oled.show()

while True:
    dist = measure_distance()
    state = move_robot(dist)

    update_display(dist, state)

    print("Distance:", dist, "| State:", state)

    time.sleep_ms(200)