#include <Adafruit_SSD1306.h> // include the Adafruit SSD1306 library
//#include <ODriveArduino.h> // include the ODrive library

#define OLED_RESET 4 // reset pin for the OLED screen
Adafruit_SSD1306 display(OLED_RESET); // create an instance of the OLED screen

#define ENCODER_A 2 // encoder pin A
#define ENCODER_B 3 // encoder pin B
#define ENCODER_BUTTON 5 // encoder button pin

#define POTENTIOMETER A0 // potentiometer pin

//#define ODRIVE_MOTOR 0 // ODrive motor number

int encoderPos = 0; // variable to store encoder position
int encoderButtonState = 0; // variable to store encoder button state
int potentiometerValue = 0; // variable to store potentiometer value

//ODriveArduino odrive(Serial); // create an instance of the ODrive

void setup() {
  pinMode(ENCODER_A, INPUT); // set encoder pin A as input
  pinMode(ENCODER_B, INPUT); // set encoder pin B as input
  pinMode(ENCODER_BUTTON, INPUT); // set encoder button pin as input
  pinMode(POTENTIOMETER, INPUT); // set potentiometer pin as input
  
  display.begin(SSD1306_SWITCHCAPVCC, 0x3C); // initialize the OLED screen
  display.clearDisplay(); // clear the screen
  display.display(); // display the cleared screen
  
  //odrive.begin(); // initialize the ODrive
  //odrive.clearErrors(); // clear any errors on the ODrive
  //odrive.setBrakeResistor(0.0f); // set the brake resistor to 0
  
  // set the starting position as the current encoder position
  encoderPos = readEncoder();
}

void loop() {
  // read the encoder position and button state
  encoderPos = readEncoder();
  encoderButtonState = digitalRead(ENCODER_BUTTON);
  
  // read the potentiometer value
  potentiometerValue = analogRead(POTENTIOMETER);
  
  // display the main menu with dumbell and arm graphics
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(WHITE);
  display.setCursor(0, 0);
  display.println("Main Menu:");
  //display.drawBitmap(0, 10, dumbell, 16, 16, WHITE); // display dumbell graphic
  //display.drawBitmap(0, 30, arm, 16, 16, WHITE); // display arm graphic
  display.display();
  
  // if the encoder button is pressed, enter the selected mode
  if (encoderButtonState == HIGH) {
    if (encoderPos == 0) { // if the encoder is at position 0, enter constant mode
      constantMode();
    } else if (encoderPos == 1) { // if the encoder is at position 1, enter variable mode
      variableMode();
    }
  }
}

// function to read the encoder position
int readEncoder() {
  int encoderPos = 0;
  int encoderAState = digitalRead(ENCODER_A);
  int encoderBState = digitalRead(ENCODER_B);
  
  if (encoderAState == HIGH && encoderBState == LOW) {
    encoderPos++;
  } else if (encoderAState == LOW && encoderBState == HIGH) {
    encoderPos--;
  }
  
  return encoderPos;
}

// function for the constant mode
void constantMode() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(WHITE);
  display.setCursor(0, 0);
  display.println("Constant Mode:");
  display.display();
  
  // read the potentiometer value and map it to a torque value
  float torque = map(potentiometerValue, 0, 1023, 0, 1);
  
  // set the torque on the ODrive
  //odrive.SetTorque(ODRIVE_MOTOR, torque);
  
  // display the current torque value
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(WHITE);
  display.setCursor(0, 0);
  display.print("Torque: ");
  display.print(torque);
  display.display();
  
  // if the encoder button is pressed, go back to the main menu
  if (encoderButtonState == HIGH) {
    loop();
  }
}

// function for the variable mode
void variableMode() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(WHITE);
  display.setCursor(0, 0);
  display.println("Variable Mode:");
  display.display();
  
  // read the potentiometer value and
}