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

# Define pins for push buttons
green_button_pin = Pin(12, Pin.IN, Pin.PULL_UP)
blue_button_pin = Pin(13, Pin.IN, Pin.PULL_UP)
yellow_button_pin = Pin(14, Pin.IN, Pin.PULL_UP)
red_button_pin = Pin(15, Pin.IN, Pin.PULL_UP)

# Define pins for LEDs
green_led_pin = Pin(16, Pin.OUT)
red_led_pin = Pin(17, Pin.OUT)

# Define buzzer pin
buzzer_pin = Pin(26, Pin.OUT)

# Define servo motor pin
servo_pin = Pin(18)
servo_pwm = PWM(servo_pin, freq=50)

# Initialize OLED display
i2c = I2C(-1, scl=Pin(22), sda=Pin(21))
oled = SSD1306_I2C(128, 64, i2c)

# Define fluid delivery parameters
ml_per_rotation = 5
rotation_time = 1  # in seconds

# Function to move servo motor for specified rotation time
def move_servo():
    servo_pwm.duty(512)  # Duty cycle for servo motor rotation
    time.sleep(rotation_time)
    servo_pwm.duty(0)

# Function to display message on OLED display
def display_message(message):
    oled.fill(0)
    oled.text(message, 0, 0)
    oled.show()

# Function to sound buzzer
def sound_buzzer():
    buzzer_pin.on()
    time.sleep(0.5)
    buzzer_pin.off()

# Main program loop
try:
    while True:
        if green_button_pin.value() == 0:
            move_servo()
            green_led_pin.on()
            red_led_pin.off()
            display_message("5ml delivered")
            sound_buzzer()
        elif blue_button_pin.value() == 0:
            move_servo()
            green_led_pin.on()
            red_led_pin.off()
            display_message("10ml delivered")
            sound_buzzer()
        elif yellow_button_pin.value() == 0:
            move_servo()
            green_led_pin.on()
            red_led_pin.off()
            display_message("15ml delivered")
            sound_buzzer()
        elif red_button_pin.value() == 0:
            move_servo()
            green_led_pin.on()
            red_led_pin.off()
            display_message("20ml delivered")
            sound_buzzer()
        else:
            green_led_pin.off()
            red_led_pin.on()

except KeyboardInterrupt:
    oled.fill(0)
    oled.show()
    print("Program stopped.")