#include <Wire.h>
#include <Adafruit_SSD1306.h>

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// A4988 Driver Pins for Motor 1
#define STEP_PIN_1 3
#define DIR_PIN_1 4
#define ENABLE_PIN_1 5

// A4988 Driver Pins for Motor 2
#define STEP_PIN_2 6
#define DIR_PIN_2 7
#define ENABLE_PIN_2 8

// Ultrasonic Sensor Pins
#define FRONT_TRIGGER_PIN 9
#define FRONT_ECHO_PIN 10
#define LEFT_TRIGGER_PIN 11
#define LEFT_ECHO_PIN A0
#define RIGHT_TRIGGER_PIN A1
#define RIGHT_ECHO_PIN A2

// Push Button Pin
#define EMERGENCY_STOP_PIN 2
#define EMERGENCY_LED_PIN 13

// Buzzer Pin
#define BUZZER_PIN A4

// Update interval for OLED display (in milliseconds)
const unsigned long DISPLAY_UPDATE_INTERVAL = 3000; // 3 seconds
unsigned long lastDisplayUpdateTime = 0;

// Wheel and Motor Parameters
#define WHEEL_DIAMETER 7.0  // Diameter of the wheel in cm (adjust as per your actual wheel size)
#define STEPS_PER_REVOLUTION 200  // Number of steps per revolution of the stepper motor
#define DISTANCE_PER_STEP ((PI * WHEEL_DIAMETER) / STEPS_PER_REVOLUTION)  // Distance covered in one step in cm

long leftMotorSteps = 0;  // Tracks steps for the left motor
long rightMotorSteps = 0; // Tracks steps for the right motor
float totalDistanceTraveled = 0;  // Total distance traveled in cm

// Current state for display
enum State {MOVING_FORWARD, OBSTACLE_DETECTED, EMERGENCY_STOP};
State currentState = MOVING_FORWARD;

int frontDistance, leftDistance, rightDistance;

void setup() {
    // Motor control pins for Motor 1
    pinMode(STEP_PIN_1, OUTPUT);
    pinMode(DIR_PIN_1, OUTPUT);
    pinMode(ENABLE_PIN_1, OUTPUT);
    
    // Motor control pins for Motor 2
    pinMode(STEP_PIN_2, OUTPUT);
    pinMode(DIR_PIN_2, OUTPUT);
    pinMode(ENABLE_PIN_2, OUTPUT);
    
    // Ultrasonic sensor pins
    pinMode(FRONT_TRIGGER_PIN, OUTPUT);
    pinMode(FRONT_ECHO_PIN, INPUT);
    pinMode(LEFT_TRIGGER_PIN, OUTPUT);
    pinMode(LEFT_ECHO_PIN, INPUT);
    pinMode(RIGHT_TRIGGER_PIN, OUTPUT);
    pinMode(RIGHT_ECHO_PIN, INPUT);
    
    // Emergency stop pin
    pinMode(EMERGENCY_STOP_PIN, INPUT);
    pinMode(EMERGENCY_LED_PIN, OUTPUT);
    
    // Buzzer pin
    pinMode(BUZZER_PIN, OUTPUT);
    
    // OLED display initialization
    display.begin(SSD1306_SWITCHCAPVCC, 0x3C);
    display.clearDisplay();
    Serial.begin(9600);
}

void loop() {
    // Check for emergency stop
    if (digitalRead(EMERGENCY_STOP_PIN) == HIGH) {
      digitalWrite(EMERGENCY_LED_PIN, HIGH);
     
        stopMotors();
        frontDistance = checkFrontDistance();
        if (currentState != EMERGENCY_STOP) {
            currentState = EMERGENCY_STOP; // Change state to emergency stop
            displayEmergencyStop(frontDistance);
        }
      return; // Exit the loop to prevent any further actions
    }

    // Move forward
    moveForward();
    frontDistance = checkFrontDistance();
    leftDistance = checkLeftDistance();
    rightDistance = checkRightDistance();

    // Update the display if the state is not EMERGENCY_STOP
    if (currentState != MOVING_FORWARD) {
        currentState = MOVING_FORWARD; // Change state to moving forward
        displayStatus("Moving Forward", frontDistance);
    }

    // Check for obstacles
    if (frontDistance < 10 || leftDistance < 10 || rightDistance < 10) {
        tone(BUZZER_PIN, 1000); // Sound buzzer if an obstacle is detected
        
        if (currentState != OBSTACLE_DETECTED) {
            currentState = OBSTACLE_DETECTED; // Change state to obstacle detected
            displayObstacleDetected(frontDistance, leftDistance, rightDistance);
            delay(1000); // Wait for a second before taking evasive action
            reverse(); // Move backward briefly
            stopMotors();
        }
    } else {
        noTone(BUZZER_PIN); // Stop the buzzer if no obstacle is detected
    }

    // Check if it's time to update the display
    if (millis() - lastDisplayUpdateTime >= DISPLAY_UPDATE_INTERVAL) {
        lastDisplayUpdateTime = millis(); // Update the last update time
        updateDisplay(frontDistance, leftDistance, rightDistance); // Update display with current distances
    }
}

void moveForward() {
    digitalWrite(DIR_PIN_1, HIGH); // Set direction for motor 1
    digitalWrite(DIR_PIN_2, HIGH); // Set direction for motor 2
    for (int i = 0; i < 200; i++) { // Number of steps for one complete revolution
        digitalWrite(STEP_PIN_1, HIGH);
        digitalWrite(STEP_PIN_2, HIGH);
        delayMicroseconds(1000); // Speed control
        digitalWrite(STEP_PIN_1, LOW);
        digitalWrite(STEP_PIN_2, LOW);
        delayMicroseconds(1000);
        
        leftMotorSteps++;  // Increment left motor step count
        rightMotorSteps++; // Increment right motor step count
    }
    updateDistanceTraveled();  // Update distance traveled after moving forward
}

void stopMotors() {
    digitalWrite(ENABLE_PIN_1, HIGH); // Disable motor driver 1
    digitalWrite(ENABLE_PIN_2, HIGH); // Disable motor driver 2
}

void reverse() {
    digitalWrite(DIR_PIN_1, LOW); // Reverse direction for motor 1
    digitalWrite(DIR_PIN_2, LOW); // Reverse direction for motor 2
    for (int i = 0; i < 200; i++) { // Number of steps for one complete revolution
        digitalWrite(STEP_PIN_1, HIGH);
        digitalWrite(STEP_PIN_2, HIGH);
        delayMicroseconds(1000); // Speed control
        digitalWrite(STEP_PIN_1, LOW);
        digitalWrite(STEP_PIN_2, LOW);
        delayMicroseconds(1000);
        
        leftMotorSteps--;  // Decrement left motor step count (since moving backwards)
        rightMotorSteps--; // Decrement right motor step count (since moving backwards)
    }
    updateDistanceTraveled();  // Update distance traveled after moving backward
}

void updateDistanceTraveled() {
    // Calculate the average steps from both motors and convert to distance
    totalDistanceTraveled = ((leftMotorSteps + rightMotorSteps) / 2.0) * DISTANCE_PER_STEP;
}

int checkFrontDistance() {
    // Trigger the front sensor
    digitalWrite(FRONT_TRIGGER_PIN, LOW);
    delayMicroseconds(2);
    digitalWrite(FRONT_TRIGGER_PIN, HIGH);
    delayMicroseconds(10);
    digitalWrite(FRONT_TRIGGER_PIN, LOW);
    
    return pulseIn(FRONT_ECHO_PIN, HIGH) * 0.034 / 2; // Convert to cm
}

int checkLeftDistance() {
    // Trigger the left sensor
    digitalWrite(LEFT_TRIGGER_PIN, LOW);
    delayMicroseconds(2);
    digitalWrite(LEFT_TRIGGER_PIN, HIGH);
    delayMicroseconds(10);
    digitalWrite(LEFT_TRIGGER_PIN, LOW);
    
    return pulseIn(LEFT_ECHO_PIN, HIGH) * 0.034 / 2; // Convert to cm
}

int checkRightDistance() {
    // Trigger the right sensor
    digitalWrite(RIGHT_TRIGGER_PIN, LOW);
    delayMicroseconds(2);
    digitalWrite(RIGHT_TRIGGER_PIN, HIGH);
    delayMicroseconds(10);
    digitalWrite(RIGHT_TRIGGER_PIN, LOW);
    
    return pulseIn(RIGHT_ECHO_PIN, HIGH) * 0.034 / 2; // Convert to cm
}

// Function to display emergency stop message with distance on OLED
void displayEmergencyStop(int frontDistance) {
    display.clearDisplay(); // Clear the display for fresh messages
    display.setTextSize(1);      
    display.setTextColor(SSD1306_WHITE);  
    display.setCursor(0, 0);
    display.print("Emergency Stop!");
    display.setCursor(0, 10);
    display.print("Front Dist: ");
    display.print(frontDistance);
    display.print(" cm");
    display.setCursor(0, 20);
    display.print("Dist Travelled: ");
    display.print(totalDistanceTraveled);
    display.print(" cm");
    display.display(); // Update the display
}

// Function to display status messages on OLED
void displayStatus(const char* message, int frontDistance) {
    display.clearDisplay(); // Clear the display for fresh messages
    display.setTextSize(1);      
    display.setTextColor(SSD1306_WHITE);  
    display.setCursor(0, 0);
    display.print(message);
    display.setCursor(0, 10);
    display.print("Front Dist: ");
    display.print(frontDistance);
    display.print(" cm");
    display.setCursor(0, 20);
    display.print("Dist Travelled: ");
    display.print(totalDistanceTraveled);
    display.print(" cm");
    display.display(); // Update the display
}

// Function to display obstacle detected message on OLED
void displayObstacleDetected(int front, int left, int right) {
    display.clearDisplay(); // Clear the display for fresh messages
    display.setTextSize(1);      
    display.setTextColor(SSD1306_WHITE);  
    display.setCursor(0, 0);
    display.print("Obstacle Detected!");
    display.setCursor(0, 10);
    display.print("Front: ");
    display.print(front);
    display.print(" cm");
    display.setCursor(0, 20);
    display.print("Left: ");
    display.print(left);
    display.print(" cm");
    display.setCursor(0, 30);
    display.print("Right: ");
    display.print(right);
    display.print(" cm");
    display.setCursor(0, 40);
    display.print("Dist Travelled: ");
    display.print(totalDistanceTraveled);
    display.print(" cm");
    display.display(); // Update the display
}

// Function to update display with current sensor data and distance
void updateDisplay(int front, int left, int right) {
    display.clearDisplay(); // Clear the display for fresh messages
    display.setTextSize(1);      
    display.setTextColor(SSD1306_WHITE);  
    display.setCursor(0, 0);
    display.print("Moving Forward!");
    display.setCursor(0, 10);
    display.print("Front: ");
    display.print(front);
    display.print(" cm");
    display.setCursor(0, 20);
    display.print("Left: ");
    display.print(left);
    display.print(" cm");
    display.setCursor(0, 30);
    display.print("Right: ");
    display.print(right);
    display.print(" cm");
    display.setCursor(0, 40);
    display.print("Distance: ");
    display.print(totalDistanceTraveled);
    display.print(" cm");
    display.display(); // Update the display
}