# Import necessary modules from MicroPython and custom drivers
from machine import Pin, I2C # Pin for GPIO control, I2C for I2C bus communication
import time # For time-related functions like delays
from buzzer import Buzzer # Custom driver for the buzzer module
from stepper import Stepper # Custom driver for controlling stepper motors
from hcsr04 import HCSR04 # Custom driver for the HC-SR04 ultrasonic sensor
from oled import I2C as OLED_I2C # Custom driver for OLED display, aliased to avoid conflict with machine.I2C
from mpu6050 import accel # Custom driver for MPU6050 accelerometer/gyroscope, using 'accel' class

# --- LED Setup ---
# Define the GPIO pin connected to the LED
LED_PIN = 12
# Initialize the LED pin as an output
led = Pin(LED_PIN, Pin.OUT)
# Ensure the LED is off when the program starts
led.value(0) 
print(f"LED on Pin {LED_PIN} initialized.")

# --- Button Setup ---
# Define the GPIO pin connected to the button
BUTTON_PIN = 27
# Initialize the button pin as an input with an internal pull-up resistor.
# This means the pin will be HIGH by default and go LOW when the button is pressed.
button = Pin(BUTTON_PIN, Pin.IN, Pin.PULL_UP) 
print(f"Button on Pin {BUTTON_PIN} initialized with PULL_UP.")

# --- Buzzer Setup ---
# Define the GPIO pin connected to the buzzer
BUZZER_PIN = 15
# Initialize the Buzzer object with the specified pin
buzzer = Buzzer(BUZZER_PIN)
print(f"Buzzer on Pin {BUZZER_PIN} initialized.")

# --- Stepper Setup ---
# Stepper motors are used for precise movement. Two steppers are set up for differential drive.

# Left Stepper Motor Configuration
LEFT_STEP_PIN = 13  # GPIO pin for stepping pulses
LEFT_DIR_PIN = 18   # GPIO pin for direction control
LEFT_ENABLE_PIN = 5 # GPIO pin for enabling/disabling the motor driver
# Initialize the Left Stepper object
left_stepper = Stepper(LEFT_STEP_PIN, LEFT_DIR_PIN, LEFT_ENABLE_PIN)
# Set the initial direction for the left stepper (0 typically means one direction, e.g., forward)
left_stepper.set_direction(0) 
# Set the speed of the left stepper in microseconds between steps (lower value = faster)
left_stepper.set_speed_us(1000) 
print(f"Left Stepper on STEP:{LEFT_STEP_PIN}, DIR:{LEFT_DIR_PIN}, ENABLE:{LEFT_ENABLE_PIN} initialized.")

# Right Stepper Motor Configuration
RIGHT_STEP_PIN = 19  # GPIO pin for stepping pulses
RIGHT_DIR_PIN = 18   # GPIO pin for direction control (Note: same as left, check wiring if intended)
RIGHT_ENABLE_PIN = 5 # GPIO pin for enabling/disabling the motor driver (Note: same as left, check wiring if intended)
# Initialize the Right Stepper object
right_stepper = Stepper(RIGHT_STEP_PIN, RIGHT_DIR_PIN, RIGHT_ENABLE_PIN)
# Set the initial direction for the right stepper (1 typically means the opposite direction of 0)
# This is crucial for forward motion when driving two wheels.
right_stepper.set_direction(1) 
# Set the speed of the right stepper in microseconds between steps
right_stepper.set_speed_us(1000) 
print(f"Right Stepper on STEP:{RIGHT_STEP_PIN}, DIR:{RIGHT_DIR_PIN}, ENABLE:{RIGHT_ENABLE_PIN} initialized.")


# --- Ultrasonic Sensor Setup ---
# Define GPIO pins for the HC-SR04 ultrasonic sensor
ULTRASONIC_TRIGGER_PIN = 4 # Pin to send trigger pulse
ULTRASONIC_ECHO_PIN = 16   # Pin to receive echo pulse
# Initialize the HCSR04 object
ultrasonic_sensor = HCSR04(ULTRASONIC_TRIGGER_PIN, ULTRASONIC_ECHO_PIN)
print(f"Ultrasonic Sensor on TRIGGER:{ULTRASONIC_TRIGGER_PIN}, ECHO:{ULTRASONIC_ECHO_PIN} initialized.")

# --- I2C Bus Initialization (Shared for OLED and MPU6050) ---
# Define GPIO pins for the I2C (Inter-Integrated Circuit) communication bus
I2C_SCL_PIN = 22 # Serial Clock Line
I2C_SDA_PIN = 21 # Serial Data Line
# Initialize the I2C bus object
i2c = I2C(scl=Pin(I2C_SCL_PIN), sda=Pin(I2C_SDA_PIN)) 
print(f"I2C bus initialized on SDA:{I2C_SDA_PIN}, SCL:{I2C_SCL_PIN}.")

# --- OLED Setup ---
# Define parameters for the OLED (Organic Light-Emitting Diode) display
OLED_WIDTH = 128     # Width of the OLED display in pixels
OLED_HEIGHT = 64     # Height of the OLED display in pixels
OLED_I2C_ADDR = 0x3c # I2C address of the OLED display
# Initialize the OLED display object, passing the I2C bus object
oled = OLED_I2C(OLED_WIDTH, OLED_HEIGHT, i2c, OLED_I2C_ADDR, external_vcc=False)
# Clear the OLED display (fill with black) on startup
oled.fill(0) 
# Show the cleared display
oled.show()
print(f"OLED Display {OLED_WIDTH}x{OLED_HEIGHT} initialized.")

# --- MPU6050 Setup ---
# Define the I2C address for the MPU6050 accelerometer/gyroscope
MPU6050_I2C_ADDR = 0x68
# Initialize the MPU6050 accelerometer object
mpu = accel(i2c, MPU6050_I2C_ADDR) 
print(f"MPU6050 Accelerometer initialized on I2C address {hex(MPU6050_I2C_ADDR)}.")


# --- Function: wait_button_press() ---
def wait_button_press():
    """
    Blocks execution until the physical button is pressed.
    It displays a message on the OLED and waits for the button's value to go low,
    then performs debouncing to prevent false triggers.
    Assumes button is wired with PULL_UP, so button.value() is 0 when pressed.
    """
    print("Waiting for button press...")
    oled.fill(0) # Clear the OLED display
    oled.text("Press Button", 0, 0) # Display "Press Button" on the first line
    oled.text("To Start", 0, 1)     # Display "To Start" on the second line
    oled.show() # Update the OLED display
    while True: # Loop indefinitely until button is pressed
        if button.value() == 0: # Check if the button is pressed (value is LOW)
            time.sleep_ms(50) # Debounce delay: wait 50 milliseconds
            if button.value() == 0: # Re-check button state after debouncing
                print("Button pressed!")
                oled.fill(0) # Clear the OLED
                oled.text("Button Pressed!", 0, 0) # Display "Button Pressed!"
                oled.show() # Update OLED
                time.sleep(0.5) # Short delay to show the "Button Pressed!" message
                break # Exit the loop as the button has been pressed
        time.sleep_ms(10) # Small delay to avoid busy-waiting and free up CPU cycles

# --- Function: get_steps_from_distance() ---
def get_steps_from_distance(distance_cm):
    """
    Calculates the number of stepper motor steps required to move a given distance in centimeters.
    This conversion is based on the stepper motor's steps per revolution and the tire's circumference.

    Args:
        distance_cm (float): The desired distance to travel in centimeters.

    Returns:
        int: The calculated number of steps, rounded to the nearest whole number.
    """
    STEPS_PER_REVOLUTION = 200 # Number of steps a single stepper motor takes for one full revolution
    TIRE_RADIUS_CM = 3.0       # Radius of the robot's tire in centimeters (60mm diameter = 30mm radius = 3cm radius)
    
    # Calculate the circumference of the tire
    circumference_cm = 2 * 3.14159 * TIRE_RADIUS_CM
    # Calculate how many steps are needed to cover one centimeter
    steps_per_cm = STEPS_PER_REVOLUTION / circumference_cm
    # Calculate the total steps required for the given distance
    total_steps = distance_cm * steps_per_cm
    
    return round(total_steps) # Return the total steps, rounded to the nearest integer

print("\n--- Robot Control Program Started ---")

# --- Main Program Loop (While True) ---
# This is the main operational loop of the robot, continuously performing tasks.
while True: 
    # Step 2: Turn the LED off at the beginning of each cycle to indicate readiness or reset state.
    led.value(0) 
    
    # Step 3: Prepare and display a message on the OLED indicating readiness for the next command.
    oled.fill(0) # Clear the entire OLED display
    oled.text("Press button to", 0, 0) # Display message on line 0
    oled.text("start", 0, 1)           # Display message on line 1
    oled.show() # Update the OLED to show the new text
    print("Ready for next command.")

    # Step 4: Call the blocking function to wait for a button press before proceeding.
    wait_button_press() 
    
    # Step 5: Activate the buzzer once to confirm the button press and start of operation.
    buzzer.beep_once() 
    
    current_distance_cm = -1 # Initialize distance variable
    try:
        # Step 6: Read the distance from the ultrasonic sensor.
        current_distance_cm = ultrasonic_sensor.distance_cm() 
        # Validate the measured distance: if it's out of typical range (e.g., 0 or > 400cm for HC-SR04),
        # set a default distance and notify the user on OLED.
        if current_distance_cm <= 0 or current_distance_cm > 400: # Max range of HC-SR04 is typically 400cm
            print(f"Distance out of sensor range or invalid: {current_distance_cm} cm. Defaulting to 10cm.")
            current_distance_cm = 10 # Default to a small distance to ensure movement still occurs
            oled.fill(0) # Clear OLED
            oled.text("Dist. Error!", 0, 0) # Display error message
            oled.text("Defaulting 10cm", 0, 1) # Explain default action
            oled.show() # Update OLED
            time.sleep(1) # Show error message for 1 second
            
        print(f"Measured Distance: {current_distance_cm} cm")
    except OSError as e:
        # Handle potential errors during ultrasonic sensor reading (e.g., sensor not connected).
        print(f"Ultrasonic Error: {e}. Cannot determine distance.")
        oled.fill(0) # Clear OLED
        oled.text("Ultra Error!", 0, 0) # Display ultrasonic error message
        oled.text(str(e.args[0]), 0, 1) # Display specific error argument if available
        oled.show() # Update OLED
        time.sleep(2) # Show error for 2 seconds
        continue # Skip the rest of the current loop iteration and go back to the beginning

    # Step 7: Calculate the number of stepper motor steps required for the measured distance.
    calculated_steps = get_steps_from_distance(current_distance_cm) 
    
    # Step 8: Display the measured distance, calculated steps, and a "Moving..." message on the OLED.
    oled.fill(0) 
    oled.text(f"Dist: {current_distance_cm}cm", 0, 0) # Display measured distance
    oled.text(f"Steps: {calculated_steps}", 0, 1)     # Display calculated steps
    oled.text("Moving...", 0, 2)                     # Indicate robot is moving
    oled.show() 
    print(f"Calculated steps for {current_distance_cm} cm: {calculated_steps} steps")

    reached = True # Step 9: Initialize a flag to track if the robot successfully reaches its target.
                   # Assume true initially, set to false if a tilt is detected.

    # Enable both stepper motors before starting the movement loop.
    left_stepper.enable() 
    right_stepper.enable()
    
    # Step 10: Loop for the calculated number of steps to move the robot.
    for step_count in range(calculated_steps): 
        # Step 11: Move both left and right motors one step.
        left_stepper.move_one_step() 
        right_stepper.move_one_step() 

        try:
            # Get accelerometer values from the MPU6050.
            mpu_values = mpu.get_values()
            # Step 12: Extract the Y-axis acceleration value. This is often used to detect forward/backward tilt.
            acY_value = mpu_values["AcY"] 
            # print(f"Step {step_count+1}/{calculated_steps}, AcY: {acY_value}") # Optional: detailed print for debugging
            
            # Step 13: Check for tilt. If the absolute AcY value exceeds a threshold, a tilt is detected.
            # The threshold (12000 in this case) depends on sensor calibration and orientation.
            if acY_value > 12000 or acY_value < -12000:
                print(f"!!! TILT DETECTED at step {step_count+1}! AcY: {acY_value}")
                reached = False # Step 14: Set the 'reached' flag to False as a tilt occurred.
                break # Step 14: Exit the movement loop immediately if a tilt is detected.
            
        except Exception as e:
            # Handle any errors that occur while reading the MPU6050 during movement.
            print(f"Error reading MPU6050 during movement: {e}")
            # Depending on robustness requirements, you might want to stop the robot (set reached=False, break)
            # if the MPU6050 fails during movement. For now, it just prints and continues.

    # Disable both stepper motors after the movement loop finishes, whether it completed
    # successfully or was interrupted by a tilt.
    left_stepper.disable() 
    right_stepper.disable()

    # Step 15: After the movement loop, check the 'reached' flag to determine the outcome.
    if reached:
        # If 'reached' is True, the robot successfully completed its movement.
        print("Robot REACHED target distance!")
        oled.fill(0) # Clear OLED
        oled.text("REACHED!", 0, 0) # Display success message
        oled.show() 
        buzzer.beep_once() # Beep once to indicate success
    else:
        # If 'reached' is False, a tilt was detected during movement.
        print("Robot TILTED! Could not reach target.")
        led.value(1) # Turn ON the LED to indicate an error/tilt
        oled.fill(0) # Clear OLED
        oled.text("TILTED!", 0, 0) # Display tilt message
        oled.show() 
        # Beep the buzzer 3 times to signal a tilt error
        for _ in range(3):
            buzzer.on()
            time.sleep(0.1) # Buzzer on for 100ms
            buzzer.off()
            time.sleep(0.1) # Buzzer off for 100ms
    
    time.sleep(3) # Step 16: Pause for 3 seconds to allow viewing of the final status message on OLED.
A4988
A4988