#include <SoftwareSerial.h>

// 4-axis stepper controller
const int NUM_MOTORS = 4;
const int MAX_SAVED_POSITIONS = 10;
const int STEP_DELAY_MICROS = 1;

// Pin setup
const int stepPins[NUM_MOTORS] = {2, 4, 6, 8};
const int dirPins[NUM_MOTORS] = {3, 5, 7, 9};
const int enablePins[NUM_MOTORS] = {A0, A1, A2, A3};

// Variables
long savedPositions[MAX_SAVED_POSITIONS][NUM_MOTORS];
long currentPos[NUM_MOTORS] = {0, 0, 0, 0};
bool isRunning = false;
bool shouldLoop = false;
bool stopNow = false;
int numSaved = 0;
int playIndex = 0;
int motorSpeeds[NUM_MOTORS] = {500, 500, 500, 500};

SoftwareSerial bluetooth(10, 11);

const int MICROSTEPS = 1;
const int STEPS_PER_REVOLUTION = 200 * MICROSTEPS;

void setup() {
  Serial.begin(9600);
  bluetooth.begin(9600);
  
  // Setup pins
  for (int i = 0; i < NUM_MOTORS; i++) {
    pinMode(stepPins[i], OUTPUT);
    pinMode(dirPins[i], OUTPUT);
    pinMode(enablePins[i], OUTPUT);
    digitalWrite(enablePins[i], LOW); // Enable motors
  }
  
  Serial.println("Stepper controller ready");
}

void loop() {
  // Check bluetooth
  if (bluetooth.available()) {
    String cmd = bluetooth.readStringUntil('\n');
    cmd.trim();
    handleCommand(cmd);
  }
  
  // Check serial
  if (Serial.available()) {
    String cmd = Serial.readStringUntil('\n');
    cmd.trim();
    handleCommand(cmd);
  }
 
  // Handle playback
  if (isRunning && !stopNow) {
    if (shouldLoop) {
      loopPlayback();
    } else {
      playNext();
    }
  }
}

void handleCommand(String cmd) {
  if (cmd.equalsIgnoreCase("S")) {
    savePosition();
  } 
  else if (cmd.equalsIgnoreCase("P")) {
    startPlayback();
  } 
  else if (cmd.equalsIgnoreCase("R")) {
    resetEverything();
  } 
  else if (cmd.equalsIgnoreCase("St")) {
    stopPlayback();
  } 
  else if (cmd.equalsIgnoreCase("LoopON")) {
    shouldLoop = true;
  } 
  else if (cmd.equalsIgnoreCase("LoopOFF")) {
    shouldLoop = false;
  } 
  else if (cmd.equalsIgnoreCase("TEST")) {
    testAllMotors();
  } 
  else if (cmd.indexOf(',') > 0) {
    parseMotorCommand(cmd);
  }
}

// Test motors
void testAllMotors() {
  Serial.println("Testing motors...");
  
  for(int motor = 0; motor < NUM_MOTORS; motor++) {
    digitalWrite(enablePins[motor], LOW);
    
    // Clockwise
    digitalWrite(dirPins[motor], HIGH);
    for (int step = 0; step < 200; step++) {
      digitalWrite(stepPins[motor], HIGH);
      delayMicroseconds(1000);
      digitalWrite(stepPins[motor], LOW);
      delayMicroseconds(1000);
    }
    
    delay(500);
    
    // Counter-clockwise
    digitalWrite(dirPins[motor], LOW);
    for (int step = 0; step < 200; step++) {
      digitalWrite(stepPins[motor], HIGH);
      delayMicroseconds(1000);
      digitalWrite(stepPins[motor], LOW);
      delayMicroseconds(1000);
    }
    
    delay(1000);
  }
}

void parseMotorCommand(String command) {
  int comma = command.indexOf(',');
  if (comma <= 0) return;
  
  int motorNum = command.substring(0, comma).toInt() - 1;
  
  if (motorNum < 0 || motorNum >= NUM_MOTORS) return;
  
  String value = command.substring(comma + 1);
  
  if (value.startsWith("SPD:")) {
    // Speed command
    int speed = value.substring(4).toInt();
    if (speed > 0) {
      motorSpeeds[motorNum] = speed;
    }
  } else {
    // Position command
    float angle = value.toFloat();
    long targetSteps = degreesToSteps(angle);
    moveMotorTo(motorNum, targetSteps);
  }
}

// Convert degrees to steps
long degreesToSteps(float degrees) {
  return (long)((STEPS_PER_REVOLUTION * degrees) / (360.0 * 16));
}

float stepsToDegrees(long steps) {
  return (steps * 360.0) / STEPS_PER_REVOLUTION;
}

// Move motor to position
void moveMotorTo(int motor, long targetSteps) {
  long stepsNeeded = targetSteps - currentPos[motor];
  
  if (stepsNeeded == 0) return;
  
  bool clockwise = (stepsNeeded > 0);
  digitalWrite(dirPins[motor], clockwise ? HIGH : LOW);
  
  int delayMicros = 1000000 / motorSpeeds[motor];
  stepsNeeded = abs(stepsNeeded);
  
  digitalWrite(enablePins[motor], LOW);
  
  for (long i = 0; i < stepsNeeded; i++) {
    digitalWrite(stepPins[motor], HIGH);
    delayMicroseconds(2);
    digitalWrite(stepPins[motor], LOW);
    delayMicroseconds(delayMicros - 2);
  }
  
  if (clockwise) {
    currentPos[motor] += stepsNeeded;
  } else {
    currentPos[motor] -= stepsNeeded;
  }
}

void savePosition() {
  if (numSaved >= MAX_SAVED_POSITIONS) return;
  
  for (int i = 0; i < NUM_MOTORS; i++) {
    savedPositions[numSaved][i] = currentPos[i];
  }
  
  numSaved++;
  Serial.println("Position saved");
}

void startPlayback() {
  if (numSaved == 0) return;

  isRunning = true;
  stopNow = false;
  playIndex = 0;
}

// Play next position
void playNext() {
  if (playIndex >= numSaved) {
    isRunning = false;
    return;
  }

  // Set directions
  for (int i = 0; i < NUM_MOTORS; i++) {
    long steps = savedPositions[playIndex][i] - currentPos[i];
    digitalWrite(dirPins[i], (steps > 0) ? HIGH : LOW);
  }
  
  // Find max steps needed
  long maxSteps = 0;
  for (int i = 0; i < NUM_MOTORS; i++) {
    long steps = abs(savedPositions[playIndex][i] - currentPos[i]);
    if (steps > maxSteps) {
      maxSteps = steps;
    }
  }
  
  // Move all motors together
  for (long step = 0; step < maxSteps; step++) {
    for (int i = 0; i < NUM_MOTORS; i++) {
      long totalSteps = abs(savedPositions[playIndex][i] - currentPos[i]);
      long stepsDone = (step * totalSteps) / maxSteps;
      long nextSteps = ((step + 1) * totalSteps) / maxSteps;
      
      if (stepsDone != nextSteps) {
        digitalWrite(stepPins[i], HIGH);
        delayMicroseconds(2);
        digitalWrite(stepPins[i], LOW);
      }
    }
    
    delayMicroseconds(STEP_DELAY_MICROS);
  }
  
  // Update positions
  for (int i = 0; i < NUM_MOTORS; i++) {
    currentPos[i] = savedPositions[playIndex][i];
  }
  
  delay(1000);
  playIndex++;
}

// Loop through all saved positions
void loopPlayback() {
  for (int pos = 0; pos < numSaved; pos++) {
    if (stopNow) break;
    
    for (int i = 0; i < NUM_MOTORS; i++) {
      long steps = savedPositions[pos][i] - currentPos[i];
      if (steps != 0) {
        digitalWrite(dirPins[i], (steps > 0) ? HIGH : LOW);
      }
    }
    
    long maxSteps = 0;
    for (int i = 0; i < NUM_MOTORS; i++) {
      long steps = abs(savedPositions[pos][i] - currentPos[i]);
      if (steps > maxSteps) {
        maxSteps = steps;
      }
    }
    
    for (long step = 0; step < maxSteps; step++) {
      for (int i = 0; i < NUM_MOTORS; i++) {
        long totalSteps = abs(savedPositions[pos][i] - currentPos[i]);
        long stepsDone = (step * totalSteps) / maxSteps;
        long nextSteps = ((step + 1) * totalSteps) / maxSteps;
        
        if (stepsDone != nextSteps) {
          digitalWrite(stepPins[i], HIGH);
          delayMicroseconds(2);
          digitalWrite(stepPins[i], LOW);
        }
      }
      
      delayMicroseconds(STEP_DELAY_MICROS);
    }
    
    for (int i = 0; i < NUM_MOTORS; i++) {
      currentPos[i] = savedPositions[pos][i];
    }
    
    delay(1000);
  }
  
  if (shouldLoop && !stopNow) {
    loopPlayback();
  } else {
    isRunning = false;
  }
}

void stopPlayback() {
  stopNow = true;
  isRunning = false;
}

void resetEverything() {
  numSaved = 0;
  isRunning = false;
  shouldLoop = false;
  playIndex = 0;
}