import machine
import time
from machine import Pin, I2C, PWM
import ssd1306
from utime import sleep 

# OLED Setup
i2c = I2C(0, scl=Pin(22), sda=Pin(21))
oled = ssd1306.SSD1306_I2C(128, 64, i2c)

# Ultrasonic Sensor Setup
trig_501 = Pin(5, Pin.OUT)
echo_501 = Pin(18, Pin.IN)
trig_505 = Pin(19, Pin.OUT)
echo_505 = Pin(23, Pin.IN)

# LED Setup
red_led = Pin(2, Pin.OUT)
yellow_led = Pin(4, Pin.OUT)
green_led = Pin(26, Pin.OUT)

# Buzzer Setup
buzzer = PWM(Pin(14, Pin.OUT))

# Servo Motor Setup
servo = machine.PWM(Pin(13), freq=50)
servo.duty(77)  # Initial position

# Function to update OLED display
def update_display(train_501_distance, train_505_distance, status):
    oled.fill(0)
    oled.text("Train 501 Dist: {} cm".format(train_501_distance), 0, 0)
    oled.text("Train 505 Dist: {} cm".format(train_505_distance), 0, 10)
    oled.text("Status: {}".format(status), 0, 20)
    oled.show()

# Function to measure distance using ultrasonic sensor
def measure_distance(trig_pin, echo_pin):
    trig_pin.off()
    time.sleep_us(2)
    trig_pin.on()
    time.sleep_us(10)
    trig_pin.off()
    
    while echo_pin.value() == 0:
        pulse_start = time.ticks_us()
    while echo_pin.value() == 1:
        pulse_end = time.ticks_us()
        
    pulse_duration = time.ticks_diff(pulse_end, pulse_start)
    distance = (pulse_duration * 0.0343) / 2
    return distance

# Function to sound the buzzer
def sound_buzzer(duration, frequency=1703):
    buzzer.init(freq=frequency, duty=400)
    sleep(duration)
    buzzer.init(freq=1, duty=0)

# Function to blink yellow LED as warning
def blink_yellow_led(blink_count=5, blink_duration=1):
    for _ in range(blink_count):
        yellow_led.on()
        sound_buzzer(0.2)
        time.sleep(blink_duration)
        yellow_led.off()
        time.sleep(blink_duration)

while True:
    train_501_distance = measure_distance(trig_501, echo_501)
    train_505_distance = measure_distance(trig_505, echo_505)
    
    if train_501_distance <= 300 or train_505_distance <= 300:
        status = "Train Detected"
        update_display(train_501_distance, train_505_distance, status)
        
        # Yellow LED warning blink
        green_led.off()
        blink_yellow_led()

        # Close gate and sound buzzer
        servo.duty(29)  # Close gate position
        red_led.on()
        sound_buzzer(10)
    else:
        status = "No Train Detected"
        update_display(train_501_distance, train_505_distance, status)
        
        # Open gate and reset LEDs
        servo.duty(77)  # Open gate position
        red_led.off()
        yellow_led.off()
        green_led.on()
        buzzer.init(freq=1, duty=0)

    # Delay to avoid rapid looping
    time.sleep(1)