// ============================================================================
// AutoScope - Automated Microscope Slide Scanner
// ============================================================================

#include <AccelStepper.h>
#include <SafeString.h>
#include <Servo.h>
#include <loopTimer.h>
#include <BufferedOutput.h>
#include "Stepper.cpp"

// ============================================================================
// PIN DEFINITIONS
// ============================================================================

// Motor interface type for AccelStepper (4-wire stepper)
#define MotorInterfaceType 4

// X-axis stepper (step/dir driver)
#define X_STEP_PIN 2
#define X_DIR_PIN 5

// Y-axis stepper (step/dir driver)
#define Y_STEP_PIN 3
#define Y_DIR_PIN 6

// Z-axis stepper (4-wire stepper motor)
#define Z_STEP_PIN 23
#define Z_DIR_PIN 25
#define Z_STEP_PIN1 27
#define Z_STEP_PIN2 29
//                                              IN1         IN3         IN2          IN4
//                                               23          27          25           29

// Limit switches
#define Z_MAX_PIN 10
#define Z_MIN_PIN 11

// Joystick
#define SW_PIN 14
#define X_PIN A15 
#define Y_PIN A14

// ============================================================================
// CONSTANTS
// ============================================================================

const int CALIBRATION_STEP_SPEED = 100;
const int MIN_SPEED = 0; 
const int MAX_SPEED = 4000; 

const int SLOW_SPEED = 20;
const int MEDIUM_SPEED = 50;
const int FAST_SPEED = 80;

const int deadZone = 6;

// ============================================================================
// GLOBAL STATE
// ============================================================================

// Joystick state
int joyX = 0;
int joyY = 0;
bool joyPressed = false;
bool lastJoyPressed = false;

// Control mode
enum ControlMode { MODE_SERIAL, MODE_JOYSTICK };
ControlMode currentMode = MODE_SERIAL;
// Triple-click detection
const int CLICKS_REQUIRED = 3;
const unsigned long CLICK_WINDOW_MS = 5000;
unsigned long clickTimes[CLICKS_REQUIRED];
int clickIndex = 0;
int speedArr[7] = {2,5,10, 30, 50, 70, 100};

// Speed/movement settings
int currentSpeedMode = MEDIUM_SPEED;
int currentMoveDistance = 100; 

// ============================================================================
// OBJECTS
// ============================================================================

createBufferedOutput(bufferedOut, 80, DROP_UNTIL_EMPTY);

millisDelay joyStickDelay;
millisDelay stepperMotorDelay;

Stepper stepperX(X_STEP_PIN, X_DIR_PIN);
Stepper stepperY(Y_STEP_PIN, Y_DIR_PIN);
AccelStepper stepperZ(MotorInterfaceType, Z_STEP_PIN, Z_STEP_PIN1, Z_DIR_PIN, Z_STEP_PIN2);

// ============================================================================
// SETUP
// ============================================================================

void setup() {
    Serial.begin(115200);
    Serial.println("AutoScope System Initializing");

    // Configure limit switch pins
    pinMode(Z_MIN_PIN, INPUT_PULLUP);
    pinMode(Z_MAX_PIN, INPUT_PULLUP);
    pinMode(SW_PIN, INPUT_PULLUP);

    for (int i = 0; i < CLICKS_REQUIRED; i++) {
        clickTimes[i] = 0;
    }

    // Initialize X stepper
    stepperX.init();
    stepperX.setMaxSpeed(MAX_SPEED);
    stepperX.setSpeed(MEDIUM_SPEED);

    // Initialize Y stepper
    stepperY.init();
    stepperY.setMaxSpeed(MAX_SPEED);
    stepperY.setSpeed(MEDIUM_SPEED);

    // Initialize Z stepper
    stepperZ.setMaxSpeed(MAX_SPEED);
    // stepperZ.setAcceleration(200.0);
    stepperZ.setSpeed(MEDIUM_SPEED);

    // Start joystick polling timer
    joyStickDelay.start(50);
    Serial.println("MODE:SERIAL"); 
}

// ============================================================================
// MAIN LOOP
// ============================================================================

void loop() {
    bufferedOut.nextByteOut(); 
    
    checkModeToggle();
    
    if (currentMode == MODE_JOYSTICK) {
        // Joystick control mode
        bool modeChanged = updateJoystickState();
        controlSteppers(modeChanged);
    }
    char c = readChar();
    processChar(c);
    runSteppers();
}

void runSteppers() {
    loopTimer.check(bufferedOut);
    stepperX.run();
    stepperY.run();
    //Different Z due to AccelStepper lib vs custom stepper for X & Y
    if(abs(stepperZ.distanceToGo()) > 10){
        stepperZ.setSpeed((stepperZ.distanceToGo() > 0 ? 1 : -1) * (currentSpeedMode * 10));
        stepperZ.runSpeed();
        // Serial.print("Speed ");
        // Serial.print(stepperZ.speed());
        // Serial.print(" CurSpeed ");
        // Serial.print(currentSpeedMode);
        // Serial.print(" Dist ");
        // Serial.println(stepperZ.distanceToGo());
    }
}

// ============================================================================
// JOYSTICK INPUT
// ============================================================================

int mapJoystickValue(int potValue) {
    return map(potValue, 0, 1024, -100, 100);
}

float mapJoystickToSpeed(float joystickValue) {
    // Apply dead zone
    if (abs(joystickValue) <= 6) {
        return 0.0;
    }
    
    // Determine direction and work with positive values
    bool isNegative = joystickValue < 0;
    float absValue = abs(joystickValue);
    
    // Normalize to 0.0 to 1.0 (excluding dead zone)
    float normalized = (float)(absValue - 6) / (100 - 6);
    
    // Apply power curve (try values between 1.5 and 3.0)
    float exponent = 2.0;
    float curved = pow(normalized, exponent);

    // Apply direction and max speed
    if (isNegative) {
        return -curved * MAX_SPEED;
    } else {
        return curved * MAX_SPEED;
    }
}

bool updateJoystickState() {
    bool locJoyPressed = !digitalRead(SW_PIN);
    int locJoyX = mapJoystickValue(analogRead(X_PIN));
    int locJoyY = mapJoystickValue(analogRead(Y_PIN));

    // Detect mode change (button press/release)
    bool modeChanged = (locJoyPressed != joyPressed);

    // Update global state
    joyX = locJoyX;
    joyY = locJoyY;
    joyPressed = locJoyPressed;

    return modeChanged;
}

// ============================================================================
// JOYSTICK -> STEPPER CONTROL
// ============================================================================

void controlSteppers(bool modeChanged) {
    // Serial.print("Pressed: "); Serial.print(joyPressed);
    // Serial.print(" joyY: "); Serial.print(joyY);
    // Serial.print(" joyX: "); Serial.println(joyX);
    delay(2.5);
    // Handle mode transitions - stop the axes we're switching away from
    if (modeChanged) {
        if (joyPressed) {
            // Switching to Z-mode: stop X/Y
            stepperX.setSpeed(0);
            stepperY.setSpeed(0);
            stopX();
            stopY();
        } else {
            // Switching to X/Y-mode: stop Z
            stepperZ.setSpeed(0);
            stopZ();
        }
    }

    // Control the active axes
    if (joyPressed) {
        controlZAxisJoyStick();
    } else {
        controlXAxisJoyStick();
        controlYAxisJoyStick();
    }
}

void controlZAxisJoyStick() {
    if (abs(joyY) > deadZone) {
        float zSpeed = mapJoystickToSpeed((float)joyY);
        stepperZ.setSpeed(zSpeed);
        stepperZ.runSpeed();
    } else {
        stopZ();
    }
}

void controlXAxisJoyStick() {
    if (abs(joyX) > deadZone) {
        stepperX.setSpeed(mapJoystickToSpeed((float)joyX));
        Serial.println(mapJoystickToSpeed((float)joyX));
        stepperX.setDirection(joyX > 0 ? 1 : -1);
        stepperX.runSpeed();
    } else {
        stopX();
    }
}

void controlYAxisJoyStick() {
    if (abs(joyY) > deadZone) {
        stepperY.setSpeed(mapJoystickToSpeed((float)joyY));
        stepperY.setDirection(joyY > 0 ? 1 : -1);
        stepperY.runSpeed();
    } else {
        stopY();
    }
}

// ============================================================================
// STOP FUNCTIONS
// ============================================================================

void stopX() {
    stepperX.stop();
}

void stopY() {
    stepperY.stop();
}

void stopZ() {
    stepperZ.setCurrentPosition(stepperZ.currentPosition());
    stepperZ.stop();
}

void stopAllSteppers() {
    stepperX.stop();
    stepperY.stop(); 
    stepperZ.stop();
    Serial.println("STOPPED");
}

// ============================================================================
// POSITION FUNCTIONS
// ============================================================================

void getCurrentPositionCM(float &x, float &y, float &z) {
    x = stepperX.currentPosition();
    y = stepperY.currentPosition();
    z = stepperZ.currentPosition();
}

void printCurrentPosition() {
    float x, y, z;
    getCurrentPositionCM(x, y, z);
    Serial.print("Current Position - X: "); Serial.print(x);
    Serial.print(", Y: "); Serial.print(y);
    Serial.print(", Z: "); Serial.print(z); 
    Serial.println("");
}

void zeroPosition() {
    stepperX.setCurrentPosition(0);
    stepperY.setCurrentPosition(0);
    stepperZ.setCurrentPosition(0);
    Serial.println("POSITION_ZEROED");
}

// ============================================================================
// MOVEMENT FUNCTIONS
// ============================================================================

void moveDirection(int xDir, int yDir, int zDir) {
    int moveDistance = getCurrentMoveDistance();
    
    if (xDir != 0) {
        stepperX.moveTo(stepperX.currentPosition() + (xDir * moveDistance));
    }
    if (yDir != 0) {
        stepperY.moveTo(stepperY.currentPosition() + (yDir * moveDistance));
    }
    if (zDir != 0) {
        stepperZ.moveTo(stepperZ.currentPosition() + (zDir * moveDistance));
    }
}


int getCurrentMoveDistance() {
    return currentMoveDistance;
}

void controlZAxis(bool dir) {
    // stepperZ.setSpeed((dir == 1 ? 1 : -1) * 500);
    stepperZ.setCurrentPosition(0);
    stepperZ.moveTo(((dir == 1 ? 1 : -1) * 100000));
    // stepperZ.runSpeed();
}

void controlXAxis(bool dir) {
    stepperX.setSpeed(currentSpeedMode);
    stepperX.setDirection(dir == 1 ? 1 : -1);
    stepperX.runSpeed();    
}

void controlYAxis(bool dir) {
    stepperY.setSpeed(currentSpeedMode);
    stepperY.setDirection(dir == 1 ? 1 : -1);
    stepperY.runSpeed();
}


// ============================================================================
// HOMING FUNCTIONS
// ============================================================================

// ============================================================================
// SPEED CONTROL
// ============================================================================

void setSpeedMode(int speed) {
    currentSpeedMode = speed;
    stepperX.setMaxSpeed(speed);
    stepperY.setMaxSpeed(speed);
    
    if (speed == SLOW_SPEED) {
        currentMoveDistance = 50;  
        Serial.println("SPEED_SLOW");
    } else if (speed == MEDIUM_SPEED) {
        currentMoveDistance = 100; 
        Serial.println("SPEED_MEDIUM");
    } else if (speed == FAST_SPEED) {
        currentMoveDistance = 200;  
        Serial.println("SPEED_FAST");
    }
}

// ============================================================================
// STATUS & DIAGNOSTICS
// ============================================================================

void reportStatus() {
    Serial.print("STATUS:");
    Serial.print("X:"); Serial.print(stepperX.currentPosition());
    Serial.print(",Y:"); Serial.print(stepperY.currentPosition());  
    Serial.print(",Z:"); Serial.print(stepperZ.currentPosition());
    Serial.print(",RUNNING:"); 
    Serial.print(stepperX.isRunning() || stepperY.isRunning() || stepperZ.isRunning());
    Serial.println();
    Serial.print("Z Min: ");
    Serial.print(digitalRead(Z_MIN_PIN));
    Serial.print(" Z Max: ");
    Serial.print(digitalRead(Z_MAX_PIN));
    Serial.println();
}

// ============================================================================
// CALIBRATION & INITIALIZATION
// ============================================================================

void emergencyStop() {
    stepperX.stop();
    stepperY.stop();
    stepperZ.stop();
    stepperX.disableOutputs();
    stepperY.disableOutputs(); 
    stepperZ.disableOutputs();
    Serial.println("EMERGENCY_STOP");
}
void stop() {
    stepperX.stop();
    stepperY.stop();
    stepperZ.stop();
    Serial.println("STOP");
}

// ============================================================================
// SERIAL COMMAND INTERFACE
// ============================================================================

char readChar() {
    char c = 0;
    if (Serial.available()) {
        c = Serial.read();
        while (Serial.available()) {
            Serial.read();
        }
    }
    return c;
}

void processChar(char c) {
    if (c == 0) {
        return;
    }
    
    // Movement commands
    if (c == 'N') {
        controlYAxis(1);       // Move North (Y+)
    } else if (c == 'S') {
        controlYAxis(0);       // Move South (Y-)
    } else if (c == 'E') {
        controlXAxis(1);         // Move East (X+)
    } else if (c == 'W') {
        controlXAxis(0);        // Move West (X-)
    } else if (c == 'U') {
        controlZAxis(1);       // Move Up (Z+)
    } else if (c == 'D') {
        controlZAxis(0);     // Move Down (Z-)
    } else if (c == 'n') {
        stopY();
    } else if (c == 's') {
        stopY();
    } else if (c == 'e') {
        stopX();
    } else if (c == 'w') {
        stopX();
    } else if (c == 'u') {
        stopZ(); 
    } else if (c == 'd') {
        stopZ();  
    } else if (c == 'X') {
        stopAllSteppers();
    
    // Position commands
    } else if (c == 'P') {
        printCurrentPosition();
    } else if (c == 'Z') {
        zeroPosition();
    
    // Speed control
    } else if (c == 'F') {
        setSpeedMode(FAST_SPEED);
    } else if (c == 'M') {
        setSpeedMode(MEDIUM_SPEED);
    } else if (c == 'L') {
        setSpeedMode(SLOW_SPEED);
    
    } else if (c == '1') {
        setSpeedMode(speedArr[0]);
    } else if (c == '2') {
        setSpeedMode(speedArr[1]);
    } else if (c == '3') {
        setSpeedMode(speedArr[2]);
    }  else if (c == '4') {
        setSpeedMode(speedArr[3]);
    } else if (c == '5') {
        setSpeedMode(speedArr[4]);
    } else if (c == '6') {
        setSpeedMode(speedArr[5]);
    }  else if (c == '0') {
        setSpeedMode(0);

    // Scanning commands
    } else if (c == 'G') {
        // goToOrigin();
    } else if (c == 'R') {
        reportStatus();
    
    // Emergency
    } else if (c == '@') {
        emergencyStop();
    
    } else if (c == '?') {
        reportStatus();
    }else if(c == '!') {
        stop();
    // Ignore whitespace
    } else if ((c == '\n') || (c == '\r')) {
        // skip
    } else if (c == 'T') {
        toggleMode(); 
    } else {
        Serial.print("Invalid cmd:"); 
        Serial.println(c);
    }
}

// ============================================================================
// MODE TOGGLE (Triple-click detection)
// ============================================================================

void checkModeToggle() {
    bool currentPressed = !digitalRead(SW_PIN);
    
    // Detect rising edge (button just pressed)
    if (currentPressed && !lastJoyPressed) {
        unsigned long now = millis();
        
        // Record this click
        clickTimes[clickIndex] = now;
        clickIndex = (clickIndex + 1) % CLICKS_REQUIRED;
        
        // Check if we have 3 clicks within the window
        int oldestIndex = clickIndex;  // After increment, this points to oldest
        unsigned long oldestClick = clickTimes[oldestIndex];
        
        // Only check if we've had at least CLICKS_REQUIRED clicks
        // and the oldest one is within our window
        if (oldestClick > 0 && (now - oldestClick) <= CLICK_WINDOW_MS) {
            toggleMode();
            // Reset click tracking
            for (int i = 0; i < CLICKS_REQUIRED; i++) {
                clickTimes[i] = 0;
            }
            clickIndex = 0;
        }
    }
    
    lastJoyPressed = currentPressed;
}

void toggleMode() {
    if (currentMode == MODE_SERIAL) {
        currentMode = MODE_JOYSTICK;
        stopAllSteppers();
        Serial.println("MODE:JOYSTICK");
    } else {
        currentMode = MODE_SERIAL;
        stopAllSteppers();
        Serial.println("MODE:SERIAL");
    }
}