/*
  Forum: https://forum.arduino.cc/t/arduino-uno-drv8825-steppermotor-problem/1358979
  Wokwi: https://wokwi.com/projects/424150168327984129

  ec2021: Changed RESET_PIN from 1 to 11 not to interfere with Serial

*/

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

// Pin definitions
#define BUTTON_FORWARD_PIN 6
#define BUTTON_NEUTRAL_PIN 7
#define BUTTON_REVERSE_PIN 8
#define LIMIT_SWITCH_PIN 9
#define RESET_PIN 11     // ec2021 -> Changed Reset Pin from 1 to 11 not to interfere with Serial!!!

#define LED_FORWARD_PIN 3
#define LED_NEUTRAL_PIN 2
#define LED_REVERSE_PIN 10

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

// Stepper motor
#define STEP_PIN 4
#define DIR_PIN 5
AccelStepper stepper(AccelStepper::DRIVER, STEP_PIN, DIR_PIN);

// Gear positions
enum GearPosition {NEUTRAL, FORWARD, REVERSE};
long initial_homing = 0;
GearPosition currentPosition = REVERSE;
const int stepsPerRevolution = 200;
// States for LEDs and display
const int ledPins[] = {LED_FORWARD_PIN, LED_NEUTRAL_PIN, LED_REVERSE_PIN};

bool homingComplete = false;

void setup() {
  // Initialize serial monitor
  Serial.begin(115200);

  // Initialize pins
  pinMode(BUTTON_FORWARD_PIN, INPUT_PULLUP);
  pinMode(BUTTON_NEUTRAL_PIN, INPUT_PULLUP);
  pinMode(BUTTON_REVERSE_PIN, INPUT_PULLUP);
  pinMode(LIMIT_SWITCH_PIN, INPUT_PULLUP);
  pinMode(RESET_PIN, INPUT_PULLUP);


  for (int i = 0; i < 3; i++) {
    pinMode(ledPins[i], OUTPUT);
  }
  ledTest();
  // Initialize the stepper motor
  stepper.setMaxSpeed(1000);
  stepper.setAcceleration(700);

  //Serial.println("Starting homing process...");  // Homing process
  while (!digitalRead(LIMIT_SWITCH_PIN) == LOW) { // Make the Stepper move CW until the switch is deactivated
    stepper.moveTo(initial_homing);
    stepper.run();
    initial_homing--;
    delay(5);
  }

  currentPosition = 0;
  // Optionally, move to neutral position after homing
  moveStepperToPosition(NEUTRAL);

  // Wait for the motor to reach neutral position
  while (stepper.distanceToGo() != 0) {
    stepper.run();
  }

  // Print "ready neutral" after homing and reaching neutral
  //Serial.println("ready neutral");
  // Initialize OLED display
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);  // Stay in infinite loop if OLED fails to initialize
  }

  display.display();
  delay(2000); // Wait for a second for startup

  // Update display initially
  homeGearShift();
  updateDisplay();
}

void ledTest() {
  // Perform LED test by lighting up LEDs in sequence
  digitalWrite(LED_FORWARD_PIN, HIGH);
  delay(500);  // LED on for 500 ms
  digitalWrite(LED_FORWARD_PIN, LOW);

  digitalWrite(LED_NEUTRAL_PIN, HIGH);
  delay(500);  // LED on for 500 ms
  digitalWrite(LED_NEUTRAL_PIN, LOW);

  digitalWrite(LED_REVERSE_PIN, HIGH);
  delay(500);  // LED on for 500 ms
  digitalWrite(LED_REVERSE_PIN, LOW);
}

void loop() {
  // Check the limit switch to set home position if not already done
  if (digitalRead(RESET_PIN) == LOW) {
    resetHoming();
  }

  // After homing, limit switch triggers are ignored (homingComplete is true)
  if (!homingComplete) {
    shiftToNeutral();
    Serial.println("Neutral");
    return;  // Do nothing until homing is completed
  }


  // Handle button presses for shifting gears
  if (digitalRead(BUTTON_FORWARD_PIN) == LOW) {
    shiftToForward();
    Serial.println("Forward");
    delay(500);
  }
  if (digitalRead(BUTTON_NEUTRAL_PIN) == LOW) {
    shiftToNeutral();
    //stepper.run();
    Serial.println("Neutral");
    delay(500);
  }
  if (digitalRead(BUTTON_REVERSE_PIN) == LOW) {
    shiftToReverse();
    Serial.println("Reverse");
    delay(500);
  }
  // Update the stepper motor position continuously
  stepper.run();

  // Update display and LEDs based on current position
  updateDisplay();
}

void shiftToForward() {
  if (currentPosition != FORWARD) {
    moveStepperToPosition(FORWARD);
    currentPosition = FORWARD;
  }
}

void shiftToNeutral() {
  if (currentPosition != NEUTRAL) {
    moveStepperToPosition(NEUTRAL);
    currentPosition = NEUTRAL;
  }
}

void shiftToReverse() {
  if (currentPosition != REVERSE) {
    moveStepperToPosition(REVERSE);
    currentPosition = REVERSE;
  }
}

void moveStepperToPosition(GearPosition position) {
  // Define the stepper motor position for each gear
  int targetPosition = 0;
  switch (position) {
    case FORWARD:
      targetPosition = 100; // example position for Forward
      break;
    case REVERSE:
      targetPosition = -100; // example position for Reverse
      break;
    case NEUTRAL:
      targetPosition = 0; // example position for Neutral
      break;
  }

  stepper.moveTo(targetPosition);
}

void homeGearShift() {
  while (digitalRead(LIMIT_SWITCH_PIN) == LOW) {  // Limit switch is LOW (not pressed)
    stepper.moveTo(initial_homing);
    stepper.run();
    initial_homing--;
    delay(5);  // Small delay to prevent excessive speed
  }
  stepper.stop();

  //currentPosition = -100;

  moveStepperToPosition(NEUTRAL);

  while (stepper.distanceToGo() != 0) {
    stepper.run();
  }

  // Serial.println("Ready Neutral");

  homingComplete = true;
  digitalWrite(LED_NEUTRAL_PIN, HIGH);
  digitalWrite(LED_FORWARD_PIN, LOW);
  digitalWrite(LED_REVERSE_PIN, LOW);
  //digitalWrite(BUTTON_NEUTRAL_PIN, LOW);
  //display.print("Ready Neutral");

  //updateDisplay();
}

void resetHoming() {
  // Reset the homing state and allow limit switch to be used again
  homingComplete = false;
  initial_homing = -1;  // Reset the homing counter
  currentPosition = NEUTRAL;

  // Move stepper to re-home position (go to limit switch again)
  stepper.setCurrentPosition(0);  // Reset current position
  homeGearShift();
}

void updateDisplay() {
  display.clearDisplay();

  // Set text size
  display.setTextSize(3);  // Set text size to 2 (larger)

  // Set text color
  display.setTextColor(SSD1306_WHITE);

  // Get the width of the text to center it
  int16_t x1, y1;
  uint16_t textWidth = 0;

  if (homingComplete) {
    // Display gear position
    switch (currentPosition) {
      case FORWARD:
        display.getTextBounds("Forward", 0, 0, &x1, &y1, &textWidth, NULL);
        display.setCursor((SCREEN_WIDTH - textWidth) / 2, SCREEN_HEIGHT / 2 - 10);  // Center horizontally and position vertically
        display.print("Forward");
        digitalWrite(LED_FORWARD_PIN, HIGH);
        digitalWrite(LED_NEUTRAL_PIN, LOW);
        digitalWrite(LED_REVERSE_PIN, LOW);
        break;
      case REVERSE:
        display.getTextBounds("Reverse", 0, 0, &x1, &y1, &textWidth, NULL);
        display.setCursor((SCREEN_WIDTH - textWidth) / 2, SCREEN_HEIGHT / 2 - 10);  // Center horizontally and position vertically
        display.print("Reverse");
        digitalWrite(LED_FORWARD_PIN, LOW);
        digitalWrite(LED_NEUTRAL_PIN, LOW);
        digitalWrite(LED_REVERSE_PIN, HIGH);
        break;
      case NEUTRAL:
        display.getTextBounds("Neutral", 0, 0, &x1, &y1, &textWidth, NULL);
        display.setCursor((SCREEN_WIDTH - textWidth) / 2, SCREEN_HEIGHT / 2 - 10);  // Center horizontally and position vertically
        display.print("Neutral");
        digitalWrite(LED_FORWARD_PIN, LOW);
        digitalWrite(LED_NEUTRAL_PIN, HIGH);
        digitalWrite(LED_REVERSE_PIN, LOW);
        break;
    }
  }
  // Update OLED display
  display.display();
}