/* 

Hardware setup :
Connect the following lines between the ESP32S3 board SPI2 and the TMC5160 driver
https://doc.riot-os.org/group__boards__esp32s3__devkit.html
https://learn.watterott.com/silentstepstick/pinconfig/tmc5160/

  MOSI ( GPIO 11)     <=> SDI/CFG1
  MISO (GPIO 13)      <=> SDO/CFG0
  SCK (GPIO 12)       <=> SCK/CFG2
  CSO (GPIO 10)       <=> CS/CFG3
  GPIO14              <=> EN
  GND                 <=> GND
  3.3V/5V             <=> VIO (depending on the processor voltage)

The TMC5160 VM pin, GND pin is powered by external power for stepper motor
The TMC5160 M2B, M2A, M1A, M1B is connect to Nema 23 stepper motor
The TMC5160 STEP pin, DIR pin is connect with RJ45

For ESP32S3 input
- DMC channel is controlled by rotary encoder (but not send to driver, only for showing to OLED display)
- rms_current (adjustable in 50mA steps) and microsteps (value from 1-256) is controlled by rotary encoder
https://arduinogetstarted.com/tutorials/arduino-rotary-encoder
- Stealth mode is controlled by switch 
https://arduinogetstarted.com/tutorials/arduino-switch

For OLED

SCL(GPIO9)   <=> SCL
SDA(GPIO8)   <=> SDA
GND          <=> GND
3V3          <=> VCC

*/
#include <SPI.h>
#include <TMCStepper.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <EEPROM.h>

#define SCREEN_WIDTH 128 // OLED display width,  in pixels
#define SCREEN_HEIGHT 64 // OLED display height, in pixels

#define EN_PIN           14 // Enable
#define CS_PIN           10 // Chip select
#define SW_MOSI          11 // Software Master Out Slave In (MOSI)
#define SW_MISO          13 // Software Master In Slave Out (MISO)
#define SW_SCK           12 // Software Slave Clock (SCK)

#define R_SENSE 0.075f // Sense resistor value

#define ENCODER_CLK   4 // Rotary encoder clock pin
#define ENCODER_DT    5 // Rotary encoder data pin
#define ENCODER_SW    6 // Rotary encoder switch pin

// Stepper driver instance
TMC5160Stepper driver(CS_PIN, R_SENSE, SW_MOSI, SW_MISO, SW_SCK);

int switchCounter = 0;

// OLED display instance
Adafruit_SSD1306 oled(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// EEPROM addresses
#define EEPROM_ADDR_INIT_FLAG  0
#define EEPROM_ADDR_CHANNEL    10
#define EEPROM_ADDR_RMS_CURRENT 20
#define EEPROM_ADDR_MICROSTEP  30
#define EEPROM_ADDR_CHOPPER_MODE   40

#define EEPROM_INIT_FLAG 0xAB // A value to signify that EEPROM is initialized

// Rotary encoder state
volatile int counter = 0;
volatile int lastCounter = 0;

volatile int modeCounter = 0;
volatile int rmsCounter = 0;
volatile int microCounter = 0;
volatile int chanCounter = 0;

unsigned long modeLastChanged = 0;

// Allowed microsteps values
const int microsteps[] = {2, 4, 8, 16, 32, 64, 128, 256};
int currentMicrostepIndex; // Variables to store microStep, starting at 16 microsteps
int currentRmsCurrent; // Variables to store current values in mA
int currentChannel; // Variables to store channel
bool ChopperMode; // Variables to store chopper mode

// Default values
const int defaultChannel = 1;
const int defaultRmsCurrent = 1200; // in mA
const int defaultMicrostepIndex = 3; // 16 microsteps
const bool defaultChopperMode = true;

void switch_to_StealthChop() {
  driver.en_pwm_mode(true); // Toggle stealthChop
  driver.pwm_autoscale(true); // Needed for stealthChop
}

void switch_to_SpreadCycle() {
  driver.en_pwm_mode(false); 
  driver.pwm_autoscale(false);
}

void IRAM_ATTR read_encoder() {
  // Encoder interrupt routine for both pins. Updates counter
  // if they are valid and have rotated a full indent

  static uint8_t old_AB = 3;  // Lookup table index
  static int8_t encval = 0;   // Encoder value  
  static const int8_t enc_states[]  = {0,-1,1,0,1,0,0,-1,-1,0,0,1,0,1,-1,0}; // Lookup table
  static unsigned long lastInterruptTime = 0;
  unsigned long interruptTime = millis();
   
  old_AB <<=2;  // Remember previous state
   
  if (digitalRead(ENCODER_CLK)) old_AB |= 0x02; // Add current state of pin A
  if (digitalRead(ENCODER_DT)) old_AB |= 0x01; // Add current state of pin B
   
  encval += enc_states[( old_AB & 0x0f )];
 
  // Update counter if encoder has rotated a full indent, that is at least 4 steps
  if( encval > 3 ) {                                    // Four steps forward
    if (interruptTime - lastInterruptTime > 40) {       // Greater than 40 milliseconds
      counter ++;                                       // Increase by 1
    } else if (interruptTime - lastInterruptTime > 20){ // Greater than 20 milliseconds
      counter += 3;                                     // Increase by 3
    } else {                                            // Faster than 20 milliseconds
      counter += 10;                                    // Increase by 10
    }
    encval = 0;
    lastInterruptTime = millis();                       // Remember time
  }
  else if( encval < -3 ) {                              // Four steps backwards
    if (interruptTime - lastInterruptTime > 40) {       // Greater than 40 milliseconds
      counter --;                                       // Decrease by 1
    } else if (interruptTime - lastInterruptTime > 20){ // Greater than 20 milliseconds
      counter -= 3;                                     // Decrease by 3
    } else {                                            // Faster than 20 milliseconds
      counter -= 10;                                    // Decrease by 10
    }
    encval = 0;
    lastInterruptTime = millis();                       // Remember time
  }
}

void setup() {
  Serial.begin(9600);

  pinMode(EN_PIN, OUTPUT);
  pinMode(ENCODER_CLK, INPUT_PULLUP);
  pinMode(ENCODER_DT, INPUT_PULLUP);
  pinMode(ENCODER_SW, INPUT_PULLUP);
  digitalWrite(EN_PIN, LOW); // Enable driver in hardware

  // Enable SPI communication
  SPI.begin();

  // Initialize the stepper driver
  driver.begin();

  EEPROM.begin(512);

  // Check if EEPROM is initialized
  if (EEPROM.read(EEPROM_ADDR_INIT_FLAG) != EEPROM_INIT_FLAG) {
    // Initialize EEPROM with default values
    EEPROM.put(EEPROM_ADDR_INIT_FLAG, EEPROM_INIT_FLAG);
    EEPROM.put(EEPROM_ADDR_CHANNEL, defaultChannel);
    EEPROM.put(EEPROM_ADDR_RMS_CURRENT, defaultRmsCurrent); // Store in 10mA units
    EEPROM.put(EEPROM_ADDR_MICROSTEP, defaultMicrostepIndex);
    EEPROM.put(EEPROM_ADDR_CHOPPER_MODE, defaultChopperMode);
    EEPROM.commit();
  }

  // Read settings from EEPROM
  currentChannel = EEPROM.read(EEPROM_ADDR_CHANNEL);
  currentRmsCurrent = EEPROM.read(EEPROM_ADDR_RMS_CURRENT);
  currentMicrostepIndex = EEPROM.read(EEPROM_ADDR_MICROSTEP);
  ChopperMode = EEPROM.read(EEPROM_ADDR_CHOPPER_MODE);

  // Validate and clamp RMS current within safe bounds
  if (currentRmsCurrent < 50 || currentRmsCurrent > 2800) {
    currentRmsCurrent = defaultRmsCurrent;
  }

  // Validate and clamp microstep index within bounds
  if (currentMicrostepIndex < 0 || currentMicrostepIndex >= sizeof(microsteps) / sizeof(microsteps[0])) {
    currentMicrostepIndex = defaultMicrostepIndex;
  }

  // Validate and clamp channel within safe bounds
  if (currentChannel < 1 || currentChannel > 32) {
    currentChannel = defaultChannel;
  }
  
  Serial.print("EEPROM currentChannel: ");
  Serial.println(currentChannel);
  Serial.print("EEPROM currentRmsCurrent: ");
  Serial.println(currentRmsCurrent);
  Serial.print("EEPROM currentMicrostepIndex: ");
  Serial.println(currentMicrostepIndex);
  Serial.print("EEPROM ChopperMode: ");
  Serial.println(ChopperMode);

  // Use interrupts for CLK and DT pins
  attachInterrupt(digitalPinToInterrupt(ENCODER_CLK), read_encoder, CHANGE);
  attachInterrupt(digitalPinToInterrupt(ENCODER_DT), read_encoder, CHANGE);

  // Apply settings to the stepper driver
  driver.toff(5); // Enables driver in software
  driver.rms_current(currentRmsCurrent); // Set motor RMS current
  driver.microsteps(microsteps[currentMicrostepIndex]); // Set microsteps

  if(ChopperMode){
    switch_to_StealthChop();
  } else {
    switch_to_SpreadCycle();
  }

  // Initialize OLED display
  if (!oled.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    while (true);
  }

  oled.clearDisplay(); // Clear display
  oled.setTextSize(1); // Set text size
  oled.setTextColor(WHITE); // Set text color
  show_oled();
}

void loop() {
  if (digitalRead(ENCODER_SW) == LOW && millis() - modeLastChanged > 300) {
    modeLastChanged = millis();
    switchCounter++;
    if (switchCounter > 3) {
      switchCounter = 0;
    }
    show_oled();
  }

  if (counter != lastCounter) {
    int delta = counter - lastCounter;
    lastCounter = counter;

    if (switchCounter == 0) {
      microCounter += delta;
      Serial.print("MICRO COUNTER: ");
      Serial.println(microCounter);
    } else if (switchCounter == 1) {
      rmsCounter += delta;
      Serial.print("RMS COUNTER: ");
      Serial.println(rmsCounter);
    } else if (switchCounter == 2) {
      modeCounter += delta;
      Serial.print("MODE COUNTER: ");
      Serial.println(modeCounter);
    } else {
      chanCounter += delta;
      Serial.print("CHAN COUNTER: ");
      Serial.println(chanCounter);
    }
  }

  // Handle mode encoder changes
  if (modeCounter != 0) {
    ChopperMode = !ChopperMode;
    if (ChopperMode == true) {
      switch_to_StealthChop();
      Serial.println("Chopper mode is Stealth");
    } else {
      switch_to_SpreadCycle();
      Serial.println("Chopper mode is SpreadCycle");
    }
    EEPROM.put(EEPROM_ADDR_CHOPPER_MODE, ChopperMode);
    EEPROM.commit();
    
    modeCounter = 0; // Reset modeCounter
    show_oled();
  }

  // Handle RMS current encoder changes
  if (rmsCounter != 0) {
    int newRmsCurrent = constrain(currentRmsCurrent + (rmsCounter * 50), 50, 2800); // Adjust RMS current within 50-2800mA
    if (newRmsCurrent != currentRmsCurrent) { // Only update if there's a change
      currentRmsCurrent = newRmsCurrent;
      driver.rms_current(currentRmsCurrent);
      EEPROM.put(EEPROM_ADDR_RMS_CURRENT, currentRmsCurrent);
      EEPROM.commit();
    }
    rmsCounter = 0; // Reset rmsCounter
    show_oled();
  }

  // Handle microsteps encoder changes
  if (microCounter != 0) {
    int newMicrostepIndex = constrain(currentMicrostepIndex + microCounter, 0, 7); // Adjust index within valid range
    if (newMicrostepIndex != currentMicrostepIndex) { // Only update if there's a change
      currentMicrostepIndex = newMicrostepIndex;
      driver.microsteps(microsteps[currentMicrostepIndex]);
      EEPROM.put(EEPROM_ADDR_MICROSTEP, currentMicrostepIndex);
      EEPROM.commit();
    }
    microCounter = 0; // Reset microCounter
    show_oled();
  }

  // Handle channel encoder changes
  if (chanCounter != 0) {
    int newChannel = constrain(currentChannel + chanCounter, 1, 32); // Adjust channel within 1-32
    if (newChannel != currentChannel) { // Only update if there's a change
      currentChannel = newChannel;
      EEPROM.put(EEPROM_ADDR_CHANNEL, currentChannel);
      EEPROM.commit();
    }
    chanCounter = 0; // Reset chanCounter
    show_oled();
  }
}

// Function to display settings on OLED
void show_oled(){
  // Update OLED display to match the format in the provided image
  oled.clearDisplay();
  
  oled.setCursor(15, 5);
  oled.print(currentChannel);
  if(switchCounter == 3){
    oled.setCursor(5, 5);
    oled.print(">");
  }

  oled.drawLine(40, 5, 40, 60, WHITE);
  oled.setCursor(55, 5);
  oled.print("MS ");
  oled.print(microsteps[currentMicrostepIndex]);
  if(switchCounter == 0){
    oled.setCursor(45, 5);
    oled.print(">");
  }

  oled.setCursor(55, 25);
  oled.print("mA ");
  oled.print(currentRmsCurrent);
  if(switchCounter == 1){
    oled.setCursor(45, 25);
    oled.print(">");
  }

  oled.setCursor(55, 45);
  oled.print(ChopperMode ? "StealthChop" : "SpreadCycle");
  if(switchCounter == 2){
    oled.setCursor(45, 45);
    oled.print(">");
  }
  oled.display();
}