#include <TM1637Display.h>  // Include the TM1637 library
#include <AccelStepper.h>   // Include the AccelStepper library for stepper control

// Pin Definitions
#define CLK 3  // Clock pin connected to D3
#define DIO 2  // Data pin connected to D2
#define BUTTON_PLUS 5  // Pin for the + button (A6)
#define BUTTON_MINUS 4  // Pin for the - button (A7)
#define BUTTON_ENTER 6  // Pin for the Enter button (A6)
#define BUTTON_BACK 8   // Pin for the Back button
#define BUTTON_PLUS10 A0 // Pin for the +10 button
#define BUTTON_MINUS10 A1 // Pin for the -10 button
#define LED_PIN 7  // Pin for the LED to signal Stage 2

// Stepper pins for A4988 controller
#define STEP_PIN_1 12   // Step pin for stepper1
#define DIR_PIN_1 11    // Direction pin for stepper1
#define STEP_PIN_2 10   // Step pin for stepper2
#define DIR_PIN_2 9    // Direction pin for stepper2

TM1637Display display(CLK, DIO);  // Create display object with CLK and DIO pins

int counter = 0;  // Variable to keep track of the displayed number in Stage 1
int counter2 = 0; // Variable for Stage 2
int counter3 = 0; // Variable for Stage 3 (loop count)
int stage = 1;    // Variable to keep track of the current stage (1, 2, or 3)

AccelStepper stepper1(AccelStepper::DRIVER, STEP_PIN_1, DIR_PIN_1);  // Stepper 1 object
AccelStepper stepper2(AccelStepper::DRIVER, STEP_PIN_2, DIR_PIN_2);  // Stepper 2 object

void setup() {
  display.setBrightness(0x0f);  // Set brightness (0x0f is maximum)
  display.showNumberDec(counter);  // Display the initial counter value for Stage 1

  // Set up button pins as inputs with pull-up resistors
  pinMode(BUTTON_PLUS, INPUT_PULLUP);
  pinMode(BUTTON_MINUS, INPUT_PULLUP);
  pinMode(BUTTON_ENTER, INPUT_PULLUP);
  pinMode(BUTTON_BACK, INPUT_PULLUP);  // Set up the Back button
  pinMode(BUTTON_PLUS10, INPUT_PULLUP); // Set up the +10 button
  pinMode(BUTTON_MINUS10, INPUT_PULLUP); // Set up the -10 button

  // Set up LED pin as output
  pinMode(LED_PIN, OUTPUT);

  // Initialize serial monitor for debugging
  Serial.begin(9600);

  // Stepper motor initialization
  stepper1.setMaxSpeed(1000);  // Set max speed for stepper1
  stepper1.setAcceleration(500);  // Set acceleration for stepper1

  stepper2.setMaxSpeed(1000);  // Set max speed for stepper2
  stepper2.setAcceleration(500);  // Set acceleration for stepper2
}

void loop() {
  int buttonPlusState = digitalRead(BUTTON_PLUS);  // Read state of the + button
  int buttonMinusState = digitalRead(BUTTON_MINUS);  // Read state of the - button
  int buttonEnterState = digitalRead(BUTTON_ENTER);  // Read state of the Enter button
  int buttonBackState = digitalRead(BUTTON_BACK);    // Read state of the Back button
  int buttonPlus10State = digitalRead(BUTTON_PLUS10); // Read state of the +10 button
  int buttonMinus10State = digitalRead(BUTTON_MINUS10); // Read state of the -10 button

  // Debugging: Print button states
  Serial.print("Button + state: ");
  Serial.println(buttonPlusState);
  Serial.print("Button - state: ");
  Serial.println(buttonMinusState);
  Serial.print("Button Enter state: ");
  Serial.println(buttonEnterState);
  Serial.print("Button Back state: ");
  Serial.println(buttonBackState);
  Serial.print("Button +10 state: ");
  Serial.println(buttonPlus10State);
  Serial.print("Button -10 state: ");
  Serial.println(buttonMinus10State);

  // Stage 1: Adjust counter value
  if (stage == 1) {
    digitalWrite(LED_PIN, LOW);  // Turn off the LED in Stage 1

    if (buttonPlusState == LOW) {
      counter++;
      if (counter > 9999) {
        counter = 0;  // Wrap around if the counter exceeds 9999
      }
      display.showNumberDec(counter);  // Update the display
      delay(200);  // Small delay to debounce
    }

    if (buttonMinusState == LOW) {
      counter--;
      if (counter < 0) {
        counter = 9999;  // Wrap around if the counter goes below 0
      }
      display.showNumberDec(counter);  // Update the display
      delay(200);  // Small delay to debounce
    }

    if (buttonPlus10State == LOW) {
      counter += 10;  // Increase counter by 10
      if (counter > 9999) {
        counter = 0;  // Wrap around if the counter exceeds 9999
      }
      display.showNumberDec(counter);  // Update the display
      delay(200);  // Small delay to debounce
    }

    if (buttonMinus10State == LOW) {
      counter -= 10;  // Decrease counter by 10
      if (counter < 0) {
        counter = 9999;  // Wrap around if the counter goes below 0
      }
      display.showNumberDec(counter);  // Update the display
      delay(200);  // Small delay to debounce
    }

    if (buttonEnterState == LOW) {
      stage = 2;  // Move to stage 2 when Enter button is pressed
      display.showNumberDec(counter2);  // Show counter2 (Stage 2)
      delay(200);  // Small delay to debounce
    }
  }

  // Stage 2: Adjust counter2 value
  else if (stage == 2) {
    digitalWrite(LED_PIN, HIGH);  // Turn on the LED in Stage 2

    if (buttonPlusState == LOW) {
      counter2++;
      
      if (counter2 > 9999) {
        counter2 = 0;  // Wrap around if the counter exceeds 9999
      }
      display.showNumberDec(counter2);  // Update the display
      delay(200);  // Small delay to debounce
    }

    if (buttonMinusState == LOW) {
      counter2--;
     
      if (counter2 < 0) {
        counter2 = 9999;  // Wrap around if the counter goes below 0
      }
      display.showNumberDec(counter2);  // Update the display
      delay(200);  // Small delay to debounce
    }
     if (buttonPlus10State == LOW) {
      counter2 += 10;  // Increase counter by 10
      if (counter2 > 9999) {
        counter2 = 0;  // Wrap around if the counter exceeds 9999
      }
      display.showNumberDec(counter2);  // Update the display
      delay(200);  // Small delay to debounce
    }

    if (buttonMinus10State == LOW) {
      counter2 -= 10;  // Decrease counter by 10
      if (counter2 < 0) {
        counter2 = 9999;  // Wrap around if the counter goes below 0
      }
      display.showNumberDec(counter2);  // Update the display
      delay(200);  // Small delay to debounce
    }

    if (buttonEnterState == LOW) {
      stage = 3;  // Move to stage 3 when Enter button is pressed
      counter3 = counter2;  // Set the number of repetitions from Stage 2
      display.showNumberDec(counter3);  // Show counter3 (Stage 3)
      delay(200);  // Small delay to debounce
    }

    if (buttonBackState == LOW) {
      stage = 1;  // Go back to stage 1 when Back button is pressed
      display.showNumberDec(counter);  // Show the value of counter (Stage 1)
      digitalWrite(LED_PIN, LOW);  // Turn off the LED
      delay(200);  // Small delay to debounce
    }
  }

  // Stage 3: Control the stepper motors using relative movements
  else if (stage == 3) {
    digitalWrite(LED_PIN, HIGH);  // Keep the LED on for Stage 3
    for (int i = 0; i < counter3; i++) {
      // Stepper 1: Move relative steps (counter steps)
      stepper1.move(counter);  // Move stepper1 by 'counter' relative steps
      while (stepper1.distanceToGo() != 0) {
        stepper1.run();  // Run stepper1 to complete its movement
      }

      // Stepper 2: Move relative 20 steps forward
      stepper2.move(20);  // Move stepper2 20 steps forward
      while (stepper2.distanceToGo() != 0) {
        stepper2.run();  // Run stepper2 to complete its movement
      }

      // Stepper 2: Move relative 20 steps backward
      stepper2.move(-20);  // Move stepper2 20 steps backward
      while (stepper2.distanceToGo() != 0) {
        stepper2.run();  // Run stepper2 to complete its movement
      }
    }
    stage = 1;  // After completing the stepper movements, return to Stage 1
    display.showNumberDec(counter);  // Show the value of counter (Stage 1)
    digitalWrite(LED_PIN, LOW);  // Turn off the LED
    delay(200);  // Small delay to debounce
  }

  delay(50);  // Small delay for stability
}