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

// Define pins for stepper motor
#define STEP_PIN 26
#define DIR_PIN 25

// Define pins for encoder
#define ENC_CLK 32
#define ENC_DT 33
#define ENC_SW 34

// Define I2C pins for LCD (default for ESP32: SDA=21, SCL=22)
#define SDA_PIN 21
#define SCL_PIN 22

// LCD setup (I2C address 0x27)
LiquidCrystal_I2C lcd(0x27, 20, 4);

// Stepper motor setup
AccelStepper stepper(AccelStepper::DRIVER, STEP_PIN, DIR_PIN);

// Encoder variables
volatile long encoderPos = 0;
volatile bool buttonPressed = false;
long lastEncoderPos = 0;

// Encoder sensitivity: how many clicks to change a value by 1
const int encoderSensitivity = 4;
long rawEncoderCount = 0;

// Menu variables
int menuItem = 0; // 0: Count, 1: Steps, 2: Speed, 3: Start
bool editing = false;
int countValue = 5; // Default count
int stepsValue = 100; // Default steps
int speedValue = 200; // Default speed in steps/second
float angleValue = 0.0; // New variable for angle

// Interrupt handlers for encoder
void IRAM_ATTR updateEncoder() {
    static unsigned long lastInterruptTime = 0;
    unsigned long interruptTime = millis();
    if (interruptTime - lastInterruptTime > 5) { // Debounce
        int clk = digitalRead(ENC_CLK);
        int dt = digitalRead(ENC_DT);
        
        if (editing) {
            if (clk != dt) {
                rawEncoderCount++;
            } else {
                rawEncoderCount--;
            }
        } else {
            if (clk == dt) {
                rawEncoderCount++;
            } else {
                rawEncoderCount--;
            }
        }
        lastInterruptTime = interruptTime;
    }
}

void IRAM_ATTR buttonISR() {
    static unsigned long lastInterruptTime = 0;
    unsigned long interruptTime = millis();
    if (interruptTime - lastInterruptTime > 200) {
        buttonPressed = true;
        lastInterruptTime = interruptTime;
    }
}

void setup() {
    Wire.begin(SDA_PIN, SCL_PIN);
    
    lcd.init();
    lcd.backlight();
    
    pinMode(ENC_CLK, INPUT_PULLUP);
    pinMode(ENC_DT, INPUT_PULLUP);
    pinMode(ENC_SW, INPUT_PULLUP);
    
    attachInterrupt(digitalPinToInterrupt(ENC_CLK), updateEncoder, CHANGE);
    attachInterrupt(digitalPinToInterrupt(ENC_SW), buttonISR, FALLING);
    
    stepper.setAcceleration(99999);
    stepper.setMaxSpeed(speedValue);
    
    updateMenu();
}

void loop() {
    long currentEncoderPos = rawEncoderCount / encoderSensitivity;
    if (currentEncoderPos != lastEncoderPos) {
        int delta = currentEncoderPos - lastEncoderPos;
        if (editing) {
            if (menuItem == 0) {
                countValue += delta;
                if (countValue < 1) countValue = 1;
            } else if (menuItem == 1) {
                stepsValue += delta * 10;
                if (stepsValue < 10) stepsValue = 10;
            } else if (menuItem == 2) {
                speedValue += delta * 50;
                if (speedValue < 100) speedValue = 100;
            }
        } else {
            menuItem += delta;
            if (menuItem < 0) menuItem = 0;
            if (menuItem > 3) menuItem = 3;
        }
        lastEncoderPos = currentEncoderPos;
        updateMenu();
    }
    
    if (buttonPressed) {
        buttonPressed = false;
        if (menuItem == 3) {
            runOscillation();
        } else {
            editing = !editing;
        }
        updateMenu();
    }
    
    stepper.run();
}

void updateMenu() {
    lcd.clear();
    
    // Calculate angle from steps (200 steps = 360 degrees)
    angleValue = stepsValue * 1.8;
    
    // Display Count
    lcd.setCursor(1, 0);
    lcd.print("Count: ");
    lcd.print(countValue);
    if (menuItem == 0 && editing) {
      lcd.setCursor(19, 0);
      lcd.print("<");
    }

    // Display Steps & Angle
    lcd.setCursor(1, 1);
    lcd.print("Steps: ");
    lcd.print(stepsValue);
    lcd.print(" (");
    lcd.print(angleValue, 1); // Display angle with one decimal place
    lcd.print(")");
    if (menuItem == 1 && editing) {
      lcd.setCursor(19, 1);
      lcd.print("<");
    }
    
    // Display Speed
    lcd.setCursor(1, 2);
    lcd.print("Speed: ");
    lcd.print(speedValue);
    lcd.print(" st/s");
    if (menuItem == 2 && editing) {
      lcd.setCursor(19, 2);
      lcd.print("<");
    }
    
    // Display Start button
    lcd.setCursor(1, 3);
    lcd.print("Start");
    
    // Display the main menu selector ">"
    lcd.setCursor(0, menuItem);
    lcd.print(">");
}

void runOscillation() {
    lcd.clear();
    lcd.setCursor(0, 0);
    lcd.print("Running...");

    stepper.setMaxSpeed(speedValue);
    
    for (int i = 1; i <= countValue; i++) {
        stepper.moveTo(stepsValue);
        while (stepper.distanceToGo() != 0) {
            stepper.run();
            lcd.setCursor(0, 1);
            lcd.print("Iter: ");
            lcd.print(i);
            lcd.print("/");
            lcd.print(countValue);
            
            lcd.setCursor(0, 2);
            lcd.print("Step: +");
            lcd.print(stepper.currentPosition());
            lcd.print("    "); 
        }

        stepper.moveTo(0);
        while (stepper.distanceToGo() != 0) {
            stepper.run();
            lcd.setCursor(0, 1);
            lcd.print("Iter: ");
            lcd.print(i);
            lcd.print("/");
            lcd.print(countValue);
            
            lcd.setCursor(0, 2);
            lcd.print("Step: -");
            lcd.print(abs(stepper.currentPosition()));
            lcd.print("    ");
        }
    }
    
    lcd.clear();
    lcd.setCursor(0, 0);
    lcd.print("Finished!");
    lcd.setCursor(0, 1);
    lcd.print("Iterations: ");
    lcd.print(countValue);
    
    while (!buttonPressed) {
        delay(10);
    }
    buttonPressed = false;
    updateMenu();
}