import machine
import time
import rp2pio  # Import for PIO state machine control

# Define pin constants
IR_A_PIN = machine.Pin(13)
IR_B_PIN = machine.Pin(14)
BUTTON_PIN = machine.Pin(15, machine.Pin.IN, pull=machine.Pin.PULLUP)  # Active LOW button

# I2C LCD configuration (adjust based on your LCD module)
I2C_SCL = machine.Pin(2)
I2C_SDA = machine.Pin(3)
LCD_ADDR = 0x27  # I2C address (check your LCD datasheet)
LCD_COLS = 16  # Number of columns
LCD_ROWS = 2  # Number of rows

# Initialize I2C object and LCD functions (replace with specific commands for your LCD)
i2c = machine.I2C(0, scl=I2C_SCL, sda=I2C_SDA, freq=100000)  # Adjust frequency if needed
def lcd_init():
    # Implement initialization commands for your LCD
    pass
def lcd_clear():
    # Implement clear display command for your LCD
    pass
def lcd_print(row, col, message):
    # Implement character/string printing commands for your LCD
    pass

# Button debounce timer (optional, adjust delay as needed)
debounce_delay = 20  # ms

# State machine definition (using PIO for efficiency)
sm = rp2pio.StateMachineProgram(
    initial=machine.Pin.LOW,  # Start with both IR pins as LOW
    out_shift_right=False,
    in_shift_right=False,
)

# State transitions for object detection
sm.add_condition(source=IR_A_PIN, edge=machine.Pin.RISE, next_state="wait_B_low")
sm.add_condition(source=IR_B_PIN, edge=machine.Pin.FALL, next_state="wait_A_high")

# State transitions for button press and system reset
sm.add_condition(source=BUTTON_PIN, edge=machine.Pin.FALL, next_state="reset")

# State actions for recording timestamps
sm.enter_state("wait_B_low", action=lambda: time_A = time.ticks_us())
sm.enter_state("wait_A_high", action=lambda: time_B = time.ticks_us())

# Reset state action (clear timestamps and display)
sm.enter_state("reset", action=lambda: [time_A = None, time_B = None, lcd_clear()])

# Initialize state machine
sm_obj = rp2pio.StateMachine(init=sm, frequency=200000)  # Adjust frequency if needed

# Set direction of IR pins (input)
IR_A_PIN.dir(machine.Pin.IN)
IR_B_PIN.dir(machine.Pin.IN)

# Main loop
time_A = None  # Initialize timestamps as None
time_B = None
last_button_press = time.ticks_ms()  # Track time for debouncing

while True:
    sm_obj.run(1)  # Run state machine once
    now = time.ticks_ms()

    # Button press debouncing
    if BUTTON_PIN.value() == 0 and (now - last_button_press) > debounce_delay:
        last_button_press = now
        sm_obj.active(0)  # Force reset state

    # Calculate elapsed time (if both timestamps are available)
    if time_A is not None and time_B is not None:
        elapsed_time_us = time_B - time_A
        time_between_ms = elapsed_time_us // 1000
        # Display time data on LCD (example, adjust formatting based on your LCD capabilities)
        lcd_print(0, 0, "Time at A: {} ms".format(time_A // 1000))
        lcd_print(1, 0, "Time at B: {} ms".format(time_B // 1000))
        lcd_print(0, 8, "Elapsed: {} ms".format(time_between_ms))