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

// OLED display dimensions
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

// Initialize OLED
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// Define HC-SR04 pins for Arduino Mega
#define TRIG 7  // Trigger pin for HC-SR04
#define ECHO 6  // Echo pin for HC-SR04

// Define motor pins for Arduino Mega
#define DIR 12   // Direction pin for motor
#define STEP 13  // Step pin for motor
#define DER 2    // Another direction pin (if needed for another motor)
#define STIP 3   // Another step pin (if needed for another motor)
#define DIS 4    // Distance threshold value

// Motor steps
const int REV = 200;  // Number of steps for the motor to move

// Variables for distance measurement
float duration_us, distance_cm;

void setup() {
  // Start serial communication
  Serial.begin(9600);

  // Set up HC-SR04 sensor pins
  pinMode(TRIG, OUTPUT);
  pinMode(ECHO, INPUT);

  // Set up motor driver pins
  pinMode(STEP, OUTPUT);
  pinMode(DIR, OUTPUT);
  pinMode(STIP, OUTPUT);
  pinMode(DER, OUTPUT);

  // Initialize OLED
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("OLED initialization failed"));
    for (;;);  // Stay here if OLED fails
  }
  display.display();
  delay(2000);
  display.clearDisplay();
  displayMessage("System Ready");
}

void loop() {
  // Measure distance using HC-SR04 sensor
  digitalWrite(TRIG, HIGH);
  delayMicroseconds(10);
  digitalWrite(TRIG, LOW);

  duration_us = pulseIn(ECHO, HIGH);
  distance_cm = 0.017 * duration_us;  // Convert to cm

  // Display distance on OLED
  displayDistance(distance_cm);

  // Check distance and control motors
  if (distance_cm > DIS) {
    digitalWrite(DIR, HIGH);
    digitalWrite(DER, HIGH);
    moveMotor();  // Move motor forward
  } else {
    digitalWrite(DIR, LOW);
    digitalWrite(DER, LOW);
    moveMotor();  // Move motor backward
  }

  // Remove or reduce this delay to improve responsiveness
  // delay(100);  
}

void moveMotor() {
  // Move the motor for REV steps
  for (int x = 0; x < REV; x++) {
    digitalWrite(STEP, HIGH);
    digitalWrite(STIP, HIGH);
    delayMicroseconds(500);  // Adjust delay to change motor speed
    digitalWrite(STEP, LOW);
    digitalWrite(STIP, LOW);
    delayMicroseconds(500);  // Adjust delay to change motor speed
  }

  // Add a small delay after changing direction to allow motor to register direction change
  delay(10); // 10ms delay to stabilize direction change if needed
}

void displayDistance(float distance) {
  // Display the distance on the OLED screen
  display.clearDisplay();
  display.setTextSize(2);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.print("Distance:");
  display.setCursor(0, 30);
  display.print(distance);
  display.print(" cm");
  display.display();
}

void displayMessage(String message) {
  // Display an initialization message on the OLED
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.println(message);
  display.display();
}