#include <AccelStepper.h>
#include <LiquidCrystal_I2C.h>

// Define pins
#define STEP_PIN 8
#define DIR_PIN 9
#define FORWARD_BTN 6
#define BACKWARD_BTN 7
#define RESET_BTN 10
#define SPEED_BTN 11  // New pin for speed button
#define FORWARD_LED 4
#define BACKWARD_LED 5
#define RESET_LED 3

// Speed settings array
const float SPEED_OPTIONS[] = {1.0, 50.0, 100.0, 1500.0, 2000.0};
const int NUM_SPEEDS = 5;

// Motor parameters
float currentSpeed = 500.0;  // Default speed
int speedIndex = 1;          // Default speed index (500)
const float ACCELERATION = 400.0;

// Create instances
AccelStepper stepper(1, STEP_PIN, DIR_PIN);
LiquidCrystal_I2C lcd(0x27, 16, 2);

// Variables
bool lastForwardState = HIGH;
bool lastBackwardState = HIGH;
bool lastSpeedState = HIGH;   // For the speed button
bool isMoving = false;
bool movingDirection = true;  // true = forward, false = backward

void setup() {
  // Initialize stepper
  stepper.setMaxSpeed(currentSpeed);
  stepper.setAcceleration(ACCELERATION);
  
  // Initialize LCD
  lcd.init();
  lcd.backlight();
  
  // Setup pins
  pinMode(FORWARD_BTN, INPUT_PULLUP);
  pinMode(BACKWARD_BTN, INPUT_PULLUP);
  pinMode(RESET_BTN, INPUT_PULLUP);
  pinMode(SPEED_BTN, INPUT_PULLUP);  // Setup for speed button
  pinMode(FORWARD_LED, OUTPUT);
  pinMode(BACKWARD_LED, OUTPUT);
  pinMode(RESET_LED, OUTPUT);
  
  // Initial display
  updateDisplay();
}

void loop() {
  // Read buttons
  bool currentForwardState = digitalRead(FORWARD_BTN);
  bool currentBackwardState = digitalRead(BACKWARD_BTN);
  bool resetState = digitalRead(RESET_BTN);
  bool currentSpeedState = digitalRead(SPEED_BTN);
  
  // Handle reset
  if (resetState == LOW) {
    resetMotor();
    delay(200);  // Simple debounce
    return;
  }
  
  // Handle speed change
  if (currentSpeedState == LOW && lastSpeedState == HIGH) {
    changeSpeed();
    delay(200);  // Simple debounce
  }
  
  // Handle forward movement
  if (currentForwardState == LOW) {
    digitalWrite(FORWARD_LED, HIGH);
    digitalWrite(BACKWARD_LED, LOW);
    isMoving = true;
    movingDirection = true;
    updateDisplay();
  }
  // Handle backward movement
  else if (currentBackwardState == LOW) {
    digitalWrite(FORWARD_LED, LOW);
    digitalWrite(BACKWARD_LED, HIGH);
    isMoving = true;
    movingDirection = false;
    updateDisplay();
  }
  // Handle stop
  else if (lastForwardState == LOW || lastBackwardState == LOW) {
    digitalWrite(FORWARD_LED, LOW);
    digitalWrite(BACKWARD_LED, LOW);
    isMoving = false;
    updateDisplay();
  }
  
  // Apply movement based on current state
  if (isMoving) {
    if (movingDirection) {
      moveForward();
    } else {
      moveBackward();
    }
  } else {
    stepper.stop();
  }
  
  // Update button states
  lastForwardState = currentForwardState;
  lastBackwardState = currentBackwardState;
  lastSpeedState = currentSpeedState;
}

void changeSpeed() {
  // Cycle through speed options
  speedIndex = (speedIndex + 1) % NUM_SPEEDS;
  currentSpeed = SPEED_OPTIONS[speedIndex];
  
  // Update motor speed
  stepper.setMaxSpeed(currentSpeed);
  
  // Display feedback
  lcd.clear();
  lcd.print("Speed changed to:");
  lcd.setCursor(0, 1);
  lcd.print(currentSpeed);
  delay(1000);
  
  // Update normal display
  updateDisplay();
}

void moveForward() {
  // Set the current speed for continuous rotation
  stepper.setSpeed(currentSpeed);
  stepper.runSpeed();
}

void moveBackward() {
  // Set the current speed for continuous rotation (negative for backward)
  stepper.setSpeed(-currentSpeed);
  stepper.runSpeed();
}

void resetMotor() {
  // Visual indication
  digitalWrite(RESET_LED, HIGH);
  
  // Stop the motor
  isMoving = false;
  stepper.stop();
  // Reset position to zero
  stepper.setCurrentPosition(0);
  
  // Display reset message
  lcd.clear();
  lcd.print("Reset");
  delay(500);
  
  // Turn off reset LED
  digitalWrite(RESET_LED, LOW);
  
  // Update display
  updateDisplay();
}

void updateDisplay() {
  lcd.clear();
  
  // Display speed setting at the top
  lcd.setCursor(0, 0);
  lcd.print("Speed: ");
  lcd.print(currentSpeed);
  
  // Display position on the second line
  lcd.setCursor(0, 1);
  lcd.print("Pos: ");
  lcd.print(stepper.currentPosition());
}