// Define pins for stepper motor and driver
const byte DIR_PIN = 2;   // Direction pin
const byte STEP_PIN = 3;  // Step pin
const byte EN_PIN = 4;    // Enable pin

// Define pins for buttons and LEDs
const byte ON_OFF_PIN = 5;   // Button to turn on/off the motor
const byte SPEED_PIN = 6;    // Button to cycle through speeds
const byte LED_PINS[] = {7, 8, 9}; // Array of LED pins for speed indicators
const byte MOTOR_LED_PIN = 10; // LED to indicate motor state

// Define motor steps per revolution
const int SPR = 200;

// Define speeds in steps per second
const int SPEEDS[] = {50, 150, 200};

// Define variables to store current state and speed
boolean state = LOW;        // Motor state (LOW = off, HIGH = on)
byte speedIndex = 0;        // Index of current speed in SPEEDS array

void setup() {
  // Set pin modes
  pinMode(DIR_PIN, OUTPUT);
  pinMode(STEP_PIN, OUTPUT);
  pinMode(EN_PIN, OUTPUT);

  pinMode(ON_OFF_PIN, INPUT_PULLUP); // Enable internal pull-up resistor for button
  pinMode(SPEED_PIN, INPUT_PULLUP);

  for (byte i = 0; i < sizeof(LED_PINS); i++) {
    pinMode(LED_PINS[i], OUTPUT);
  }
  pinMode(MOTOR_LED_PIN, OUTPUT);

  setMotorState(false);   // Turn off motor initially
}

void loop() {
  if (digitalRead(ON_OFF_PIN) == LOW) { // If button is pressed
    state = !state; // Toggle motor state
    setMotorState(state);
    delay(50); // Debounce delay
  }

  if (state) { // If motor is on
    if (digitalRead(SPEED_PIN) == LOW) { // If button is pressed
      speedIndex = (speedIndex + 1) % (sizeof(SPEEDS)/sizeof(*SPEEDS)); // Cycle to next speed
      setSpeedIndicator(speedIndex);
      delay(50); // Debounce delay
    }

    // Rotate clockwise for one revolution at current speed
    rotateMotor(SPEEDS[speedIndex], true);
    
    // If motor is turned off during clockwise rotation, exit the loop
    if (!state) {
      return;
    }

    // Rotate counter-clockwise for one revolution at current speed
    rotateMotor(SPEEDS[speedIndex], false);

    // If motor is turned off during counter-clockwise rotation, exit the loop
    if (!state) {
      return;
    }
  }
}

void rotateMotor(int speed, boolean dir) {
  int interval = (100000000/speed)/SPR; // Calculate interval between steps in microseconds
  digitalWrite(DIR_PIN, dir);      // Set direction based on dir parameter

  for (int i = 0; i < SPR; i++) {
    digitalWrite(STEP_PIN, HIGH);
    delayMicroseconds(interval/2);
    digitalWrite(STEP_PIN, LOW);
    delayMicroseconds(interval/2);
  }
}

void setMotorState(boolean state) {
  digitalWrite(EN_PIN, !state);    // Enable/disable motor driver
  digitalWrite(MOTOR_LED_PIN, state); // Turn on/off LED to indicate motor state
  setSpeedIndicator(state ? speedIndex : 255); // Turn off all speed indicators when motor is off
}

void setSpeedIndicator(byte speedIndex) {
  for (byte i = 0; i < sizeof(LED_PINS); i++) {
    digitalWrite(LED_PINS[i], (i == speedIndex));
  }
}