#include <Arduino.h>

const int MOTOR_COUNT = 3; // We're using 3 motors (X, Y, Z)

// Pin definitions for motor control
const int EN[MOTOR_COUNT] = {13, 10, 7};    // Enable pins
const int STEP[MOTOR_COUNT] = {12, 9, 6};   // Step pins
const int DIR[MOTOR_COUNT] = {11, 8, 5};    // Direction pins

// Pin definitions for limit switches
#define LIMITX 22
#define LIMITY 24
#define LIMITZ 26

// Pin definition for stop button
#define STOP 23

// Motor movement parameters
const int stepDelay = 800;  // Microseconds between steps (controls speed)
const int homingSpeed = 1200; // Slower speed for homing (microseconds)

// Home directions (1 = positive direction, 0 = negative direction)
const bool HOME_DIR[MOTOR_COUNT] = {0, 0, 0};  // Assuming homing in negative direction

// Limit switch pins in order X, Y, Z
const int limitSwitches[MOTOR_COUNT] = {LIMITX, LIMITY, LIMITZ};

// Steps per unit (adjust based on your machine)
const float STEPS_PER_UNIT[MOTOR_COUNT] = {80.0, 80.0, 400.0}; // Example: 80 steps per mm for X/Y, 400 for Z

// Current position in steps
long currentPosition[MOTOR_COUNT] = {0, 0, 0};

// G29 probe positions
const int G29_POINTS = 7;
float G29_COORDINATES[G29_POINTS][MOTOR_COUNT] = {
  {1, 1, 1},
  {1, 1, 5},
  {30, 1, 1},
  {30, 1, 5},
  {30, 30, 1},
  {30, 30, 5},
  {1, 30, 5}
};

String inputString = "";      // String to hold incoming serial data
boolean stringComplete = false;  // Whether the string is complete

void setup() {
  Serial.begin(115200);
  inputString.reserve(50);
  
  // Initialize motor control pins
  for (int i = 0; i < MOTOR_COUNT; i++) {
    pinMode(EN[i], OUTPUT);
    pinMode(STEP[i], OUTPUT);
    pinMode(DIR[i], OUTPUT);
    
    // Enable motors (LOW = enabled for most drivers)
    digitalWrite(EN[i], LOW);
  }
  
  // Initialize limit switches with pull-up resistors
  pinMode(LIMITX, INPUT_PULLUP);
  pinMode(LIMITY, INPUT_PULLUP);
  pinMode(LIMITZ, INPUT_PULLUP);
  
  // Initialize stop button with pull-up resistor
  pinMode(STOP, INPUT_PULLUP);
  
  Serial.println("System initialized. Send G28 to home all axes or G29 to run probe sequence.");
}

void loop() {
  // Process complete commands
  if (stringComplete) {
    processCommand(inputString);
    inputString = "";
    stringComplete = false;
  }
  
  // Check for stop button press
  if (digitalRead(STOP) == LOW) {
    emergencyStop();
  }
}

void serialEvent() {
  while (Serial.available()) {
    char inChar = (char)Serial.read();
    inputString += inChar;
    
    // If the incoming character is a newline, set a flag
    if (inChar == '\n') {
      stringComplete = true;
    }
  }
}

void processCommand(String command) {
  command.trim();  // Remove any whitespace
  
  if (command.startsWith("G28")) {
    Serial.println("Executing homing sequence (G28)");
    homeAllAxes();
  } 
  else if (command.startsWith("G29")) {
    Serial.println("Executing G29 probe sequence");
    executeG29Sequence();
  }
  else {
    Serial.println("Unknown command: " + command);
  }
}

void homeAllAxes() {
  // Home in sequence: X, Y, then Z
  for (int axis = 0; axis < MOTOR_COUNT; axis++) {
    homeAxis(axis);
    delay(500); // Pause between axes
  }
  
  // After homing, reset the current position to 0 for all axes
  for (int i = 0; i < MOTOR_COUNT; i++) {
    currentPosition[i] = 0;
  }
  
  Serial.println("Homing complete - machine at home position");
}

void homeAxis(int axis) {
  char axisName;
  switch(axis) {
    case 0: axisName = 'X'; break;
    case 1: axisName = 'Y'; break;
    case 2: axisName = 'Z'; break;
  }
  
  Serial.print("Homing ");
  Serial.print(axisName);
  Serial.println(" axis...");
  
  // Set direction for homing
  digitalWrite(DIR[axis], HOME_DIR[axis]);
  
  // Move until limit switch is triggered
  while (digitalRead(limitSwitches[axis]) == HIGH) { // Assuming active LOW switches
    // Check for stop button
    if (digitalRead(STOP) == LOW) {
      emergencyStop();
      return;
    }
    
    // Step the motor
    digitalWrite(STEP[axis], HIGH);
    delayMicroseconds(homingSpeed / 2);
    digitalWrite(STEP[axis], LOW);
    delayMicroseconds(homingSpeed / 2);
  }
  
  Serial.print(axisName);
  Serial.println(" axis homed successfully");
}
void moveToPosition(float x, float y, float z) {
  // Target positions in steps
  long targetSteps[MOTOR_COUNT];
  
  // Calculate target steps from coordinates
  targetSteps[0] = (long)(x * STEPS_PER_UNIT[0]);
  targetSteps[1] = (long)(y * STEPS_PER_UNIT[1]);
  targetSteps[2] = (long)(z * STEPS_PER_UNIT[2]);
  
  // Print steps information before moving
  Serial.print("Movement to (");
  Serial.print(x);
  Serial.print(", ");
  Serial.print(y);
  Serial.print(", ");
  Serial.print(z);
  Serial.println("):");
  
  Serial.print("  X-axis steps: ");
  Serial.print(targetSteps[0] - currentPosition[0]);
  Serial.print(" (");
  Serial.print(currentPosition[0]);
  Serial.print(" -> ");
  Serial.print(targetSteps[0]);
  Serial.println(")");
  
  Serial.print("  Y-axis steps: ");
  Serial.print(targetSteps[1] - currentPosition[1]);
  Serial.print(" (");
  Serial.print(currentPosition[1]);
  Serial.print(" -> ");
  Serial.print(targetSteps[1]);
  Serial.println(")");
  
  Serial.print("  Z-axis steps: ");
  Serial.print(targetSteps[2] - currentPosition[2]);
  Serial.print(" (");
  Serial.print(currentPosition[2]);
  Serial.print(" -> ");
  Serial.print(targetSteps[2]);
  Serial.println(")");
  
  // Move each axis to its target position
  for (int axis = 0; axis < MOTOR_COUNT; axis++) {
    moveAxisToPosition(axis, targetSteps[axis]);
  }
}

void moveAxisToPosition(int axis, long targetPosition) {
  // Calculate steps to move (positive = forward, negative = backward)
  long stepsToMove = targetPosition - currentPosition[axis];
  
  if (stepsToMove == 0) return; // No movement needed
  
  char axisName = 'X' + axis; // X=0, Y=1, Z=2
  
  // Log movement start
  Serial.print("Moving ");
  Serial.print(axisName);
  Serial.print(" by ");
  Serial.print(stepsToMove);
  Serial.println(" steps");
  
  // Set direction
  boolean direction = (stepsToMove > 0);
  digitalWrite(DIR[axis], direction);
  
  // Calculate absolute steps
  long absSteps = abs(stepsToMove);
  
  // Move the required number of steps
  for (long i = 0; i < absSteps; i++) {
    // Check for stop button
    if (digitalRead(STOP) == LOW) {
      emergencyStop();
      return;
    }
    
    // Step the motor
    digitalWrite(STEP[axis], HIGH);
    delayMicroseconds(stepDelay / 2);
    digitalWrite(STEP[axis], LOW);
    delayMicroseconds(stepDelay / 2);
    
    // Update current position as we go (for more accurate tracking)
    currentPosition[axis] += direction ? 1 : -1;
    
    // Print progress every 50 steps or at the end
    if (i % 50 == 0 || i == absSteps - 1) {
      Serial.print(axisName);
      Serial.print(" position: ");
      Serial.print(currentPosition[axis]);
      Serial.print(" / ");
      Serial.print(targetPosition);
      Serial.print(" (");
      Serial.print((i+1));
      Serial.print("/");
      Serial.print(absSteps);
      Serial.println(" steps)");
    }
  }
  
  // Final position confirmation
  Serial.print(axisName);
  Serial.print("-axis movement complete. Position: ");
  Serial.println(currentPosition[axis]);
}

// You may also want to enhance the G29 sequence logging
void executeG29Sequence() {
  Serial.println("Starting G29 probe sequence");
  
  // Move through each probe point
  for (int point = 0; point < G29_POINTS; point++) {
    Serial.print("\n>> G29 Point ");
    Serial.print(point + 1);
    Serial.print(" of ");
    Serial.print(G29_POINTS);
    Serial.print(": Target X");
    Serial.print(G29_COORDINATES[point][0]);
    Serial.print(" Y");
    Serial.print(G29_COORDINATES[point][1]);
    Serial.print(" Z");
    Serial.println(G29_COORDINATES[point][2]);
    
    moveToPosition(G29_COORDINATES[point][0], G29_COORDINATES[point][1], G29_COORDINATES[point][2]);
    
    Serial.print("Arrived at point ");
    Serial.print(point + 1);
    Serial.print(": Actual X");
    Serial.print((float)currentPosition[0]/STEPS_PER_UNIT[0]);
    Serial.print(" Y");
    Serial.print((float)currentPosition[1]/STEPS_PER_UNIT[1]);
    Serial.print(" Z");
    Serial.println((float)currentPosition[2]/STEPS_PER_UNIT[2]);
    
    // Pause briefly at each point (for readings or to simulate probe action)
    delay(500);
  }
  
  Serial.println("\nG29 probe sequence completed");
}

void emergencyStop() {
  Serial.println("EMERGENCY STOP ACTIVATED");
  
  // Disable all motors
  for (int i = 0; i < MOTOR_COUNT; i++) {
    digitalWrite(EN[i], HIGH); // HIGH = disabled for most drivers
  }
  
  // Wait until stop button is released
  while (digitalRead(STOP) == LOW) {
    delay(100);
  }
  
  Serial.println("Emergency stop cleared. Re-enabling motors.");
  
  // Re-enable motors
  for (int i = 0; i < MOTOR_COUNT; i++) {
    digitalWrite(EN[i], LOW);
  }
}