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

// LCD initialization with I2C address 0x27, 20x4
LiquidCrystal_I2C lcd(0x27, 20, 4);

// Joystick pins
const int JOY_X = A0;
const int JOY_Y = A1;
const int JOY_SW = 0;  // D0

// Stepper definitions
const int stepPin = 2;    // D2
const int dirPin = 3;     // D3
const int enablePin = 1;  // D1
AccelStepper stepper(AccelStepper::DRIVER, stepPin, dirPin);

// Relay pin
const int relayPin = 4;

// Constants
const int STEPS_PER_MM = 200;     // 200 steps * 2 (half-step) / 2mm pitch
const long MAX_SPEED = 1000;      // Normal speed for sequence
const long MANUAL_SPEED = 8000;   // Much higher speed for manual control
const long POSITIONING_SPEED = 1000; // Medium speed for position setting
const long ACCELERATION = 2000;    // Erhöhte Beschleunigung für schnelleres Anhalten
const float BACKLASH = 2.0;       // Backlash compensation in mm

// Joystick deadzone
const int DEADZONE_LOW = 400;     // Untere Grenze der Deadzone
const int DEADZONE_HIGH = 600;    // Obere Grenze der Deadzone

// Menu variables
int currentMenu = 0;
const int MENU_ITEMS = 4;
String menuItems[MENU_ITEMS] = {"Manuell", "Anfahren", "Aufnahmen", "Start"};

// Position variables
long startPosition = 0;
long endPosition = 0;
int numShots = 0;
bool isRunning = false;

// Function to enable/disable motor
void enableMotor(bool enable) {
  if(enable) {
    stepper.enableOutputs();
  } else {
    stepper.disableOutputs();
  }
}

// Function to convert steps to mm with 2 decimal places
String stepsToMM(long steps) {
  char buf[10];
  dtostrf(float(steps) / STEPS_PER_MM, 6, 2, buf);
  return String(buf);
}

// Function to handle backlash compensation
void gotoPositionWithBacklash(long targetPosition) {
  enableMotor(true);  // Enable motor before movement
  long currentPos = stepper.currentPosition();
  if (targetPosition < currentPos) {
    // When moving backwards, overshoot and then approach from behind
    long overshootPosition = targetPosition - (BACKLASH * STEPS_PER_MM);
    stepper.moveTo(overshootPosition);
    while (stepper.distanceToGo() != 0) {
      stepper.run();
    }
    stepper.moveTo(targetPosition);
    while (stepper.distanceToGo() != 0) {
      stepper.run();
    }
  } else {
    // When moving forwards, go directly to position
    stepper.moveTo(targetPosition);
    while (stepper.distanceToGo() != 0) {
      stepper.run();
    }
  }
  enableMotor(false);  // Disable motor after movement
}

void setup() {
  // Initialize LCD
  lcd.init();
  lcd.backlight();
  
  // Initialize stepper
  stepper.setMaxSpeed(MAX_SPEED);
  stepper.setAcceleration(ACCELERATION);
  stepper.setEnablePin(enablePin);
  stepper.setPinsInverted(false, false, true);
  enableMotor(false);  // Start with motor disabled
  
  // Initialize relay
  pinMode(relayPin, OUTPUT);
  digitalWrite(relayPin, LOW);
  
  // Initialize joystick button
  pinMode(JOY_SW, INPUT_PULLUP);
  
  displayMenu();
}

void loop() {
  if (isRunning) {
    runSequence();
  } else {
    handleMenu();
  }
  stepper.run();
}

void displayMenu() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("=== Makro Slider ===");
  
  // Display up to 3 menu items
  for(int i = 0; i < 3 && i + currentMenu < MENU_ITEMS; i++) {
    lcd.setCursor(0, i + 1);
    lcd.print(currentMenu + i == currentMenu ? "> " : "  ");
    lcd.print(menuItems[currentMenu + i]);
  }
}

void handleMenu() {
  int y = analogRead(JOY_Y);
  
  // Menu navigation
  if (y > 800 && currentMenu < MENU_ITEMS - 1) {
    currentMenu++;
    displayMenu();
    delay(200);
  } else if (y < 200 && currentMenu > 0) {
    currentMenu--;
    displayMenu();
    delay(200);
  }
  
  // Menu selection
  if (digitalRead(JOY_SW) == LOW) {
    delay(50);  // Debounce
    if (digitalRead(JOY_SW) == LOW) {
      executeMenuItem(currentMenu);
    }
  }
}

void executeMenuItem(int menu) {
  switch (menu) {
    case 0:  // Manual
      manualControl();
      break;
    case 1:  // Set positions
      setPositions();
      break;
    case 2:  // Set number of shots
      setNumShots();
      break;
    case 3:  // Start
      if (validateParameters()) {
        isRunning = true;
        startSequence();
      }
      break;
  }
}

void manualControl() {
  lcd.clear();
  lcd.print("=== Manual Control ===");
  lcd.setCursor(0, 3);
  lcd.print("Press SW to confirm");
  
  // Set high speed for manual control
  long originalSpeed = stepper.maxSpeed();
  long originalAccel = stepper.acceleration();
  stepper.setMaxSpeed(MANUAL_SPEED);
  stepper.setAcceleration(ACCELERATION);
  
  while (digitalRead(JOY_SW) == HIGH) {
    int x = analogRead(JOY_X);
    
    // Implement deadzone and motor control
    if (x > DEADZONE_HIGH || x < DEADZONE_LOW) {
      enableMotor(true);
      if (x > DEADZONE_HIGH) {
        stepper.move(STEPS_PER_MM);
      } else {
        stepper.move(-STEPS_PER_MM);
      }
    } else {
      stepper.stop();
      stepper.setCurrentPosition(stepper.currentPosition());
      enableMotor(false);
    }
    
    stepper.run();
  }
  
  // Reset speed and acceleration to normal
  stepper.setMaxSpeed(originalSpeed);
  stepper.setAcceleration(originalAccel);
  enableMotor(false);
  delay(200);
  displayMenu();
}

void setPositions() {
  // Set start position
  lcd.clear();
  lcd.print("=== Set Start Pos ===");
  lcd.setCursor(0, 1);
  lcd.print("Current: ");
  lcd.setCursor(0, 3);
  lcd.print("Press SW to confirm");
  
  // Set medium speed for positioning
  long originalSpeed = stepper.maxSpeed();
  stepper.setMaxSpeed(POSITIONING_SPEED);
  
  while (digitalRead(JOY_SW) == HIGH) {
    manualPositioning();
  }
  
  // Set this position as zero
  stepper.setCurrentPosition(0);
  startPosition = 0;
  delay(200);
  
  // Set end position
  lcd.clear();
  lcd.print("=== Set End Pos ===");
  lcd.setCursor(0, 1);
  lcd.print("Current: ");
  lcd.setCursor(0, 3);
  lcd.print("Press SW to confirm");
  
  while (digitalRead(JOY_SW) == HIGH) {
    manualPositioning();
  }
  endPosition = stepper.currentPosition();
  
  // Reset speed to normal and disable motor
  stepper.setMaxSpeed(originalSpeed);
  enableMotor(false);
  delay(200);
  
  displayMenu();
}

void manualPositioning() {
  int x = analogRead(JOY_X);
  
  // Implement deadzone and motor control
  if (x > DEADZONE_HIGH || x < DEADZONE_LOW) {
    enableMotor(true);
    if (x > DEADZONE_HIGH) {
      stepper.move(STEPS_PER_MM / 5);
    } else {
      stepper.move(-STEPS_PER_MM / 5);
    }
  } else {
    stepper.stop();
    stepper.setCurrentPosition(stepper.currentPosition());
    enableMotor(false);
  }
  
  stepper.run();
  
  lcd.setCursor(9, 1);
  lcd.print(stepsToMM(stepper.currentPosition()));
  lcd.print("mm    ");
}

void setNumShots() {
  lcd.clear();
  lcd.print("=== Set # of Shots ===");
  
  while (digitalRead(JOY_SW) == HIGH) {
    int y = analogRead(JOY_Y);
    if (y > 800 && numShots < 100) {
      numShots++;
    } else if (y < 200 && numShots > 2) {
      numShots--;
    }
    
    // Calculate and display step distance
    float stepDistance = 0;
    if (numShots > 1) {
      stepDistance = float(endPosition - startPosition) / ((numShots - 1) * STEPS_PER_MM);
    }
    
    lcd.setCursor(0, 1);
    lcd.print("Shots: ");
    lcd.print(numShots);
    lcd.print("   ");
    
    lcd.setCursor(0, 2);
    lcd.print("Step dist: ");
    char buf[10];
    dtostrf(stepDistance, 6, 2, buf);
    lcd.print(buf);
    lcd.print("mm   ");
    
    delay(100);
  }
  delay(200);
  displayMenu();
}

bool validateParameters() {
  if (numShots < 2) {
    lcd.clear();
    lcd.print("Error!");
    lcd.setCursor(0, 1);
    lcd.print("Set # of shots!");
    lcd.setCursor(0, 3);
    lcd.print("Press SW to continue");
    while(digitalRead(JOY_SW) == HIGH) {}
    delay(200);
    return false;
  }
  if (startPosition == endPosition) {
    lcd.clear();
    lcd.print("Error!");
    lcd.setCursor(0, 1);
    lcd.print("Set positions!");
    lcd.setCursor(0, 3);
    lcd.print("Press SW to continue");
    while(digitalRead(JOY_SW) == HIGH) {}
    delay(200);
    return false;
  }
  return true;
}

void startSequence() {
  long stepDistance = (endPosition - startPosition) / (numShots - 1);
  
  lcd.clear();
  lcd.print("=== Running... ===");
  lcd.setCursor(0, 3);
  lcd.print("SW = Cancel");
  
  gotoPositionWithBacklash(startPosition);
  
  for (int i = 0; i < numShots; i++) {
    if (digitalRead(JOY_SW) == LOW) {
      lcd.clear();
      lcd.print("=== Cancelled ===");
      lcd.setCursor(0, 1);
      lcd.print("Returning to start...");
      delay(1000);
      gotoPositionWithBacklash(startPosition);
      isRunning = false;
      delay(1000);
      displayMenu();
      return;
    }
    
    // Move to position first
    long targetPos = startPosition + (i * stepDistance);
    gotoPositionWithBacklash(targetPos);
    
    // Then update display
    lcd.setCursor(0, 1);
    lcd.print("Shot: ");
    lcd.print(i + 1);
    lcd.print("/");
    lcd.print(numShots);
    lcd.print("   ");
    
    lcd.setCursor(0, 2);
    lcd.print("Position: ");
    lcd.print(stepsToMM(stepper.currentPosition()));
    lcd.print("mm   ");
    
    // Take photo sequence
    lcd.setCursor(0, 3);
    lcd.print("Wait before shot...");
    delay(3000);
    
    lcd.setCursor(0, 3);
    lcd.print("Taking photo...   ");
    digitalWrite(relayPin, HIGH);
    delay(1000);
    digitalWrite(relayPin, LOW);
    
    lcd.setCursor(0, 3);
    lcd.print("Wait after shot...");
    delay(3000);
  }
  
  lcd.setCursor(0, 3);
  lcd.print("Returning...      ");
  gotoPositionWithBacklash(startPosition);
  
  isRunning = false;
  displayMenu();
}
void runSequence() {
  if (!stepper.isRunning() && !isRunning) {
    displayMenu();
  }
}