#include <Wire.h>
#include <Adafruit_SSD1306.h>
#include <Adafruit_MPU6050.h>

// Pin Definitions
#define BUTTON1_PIN 2
#define BUTTON2_PIN 3
#define POT1_PIN A0
#define POT2_PIN A1

// OLED Setup
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 32
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// MPU6050 Setup
Adafruit_MPU6050 mpu;

// Variables
int state = 0;          // Current state (0: Pot1, 1: Pot2, 2: MPU6050)
int mpuAxis = 0;        // MPU6050 axis selector (0: Acc X, 1: Acc Y, ..., 5: Gyro Z)

// Debouncing variables
unsigned long lastDebounceTime1 = 0;
unsigned long lastDebounceTime2 = 0;
const unsigned long debounceDelay = 120;

// Setup
void setup() {
  pinMode(BUTTON1_PIN, INPUT_PULLUP);
  pinMode(BUTTON2_PIN, INPUT_PULLUP);
  Serial.begin(115200);

  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println("OLED initialization failed.");
    while (true);
  }

  display.clearDisplay();
  display.display();
  display.setTextSize(1);
  display.setTextColor(WHITE);

  if (!mpu.begin()) {
    Serial.println("MPU6050 initialization failed.");
    while (true);
  }
}

// Reads the selected sensor
float readSensor() {
  if (state == 0) return analogRead(POT1_PIN);
  if (state == 1) return analogRead(POT2_PIN);

  // MPU6050 readings
  sensors_event_t a, g, temp;
  mpu.getEvent(&a, &g, &temp);
  if (mpuAxis == 0) return a.acceleration.x;
  if (mpuAxis == 1) return a.acceleration.y;
  if (mpuAxis == 2) return a.acceleration.z;
  if (mpuAxis == 3) return g.gyro.x;
  if (mpuAxis == 4) return g.gyro.y;
  if (mpuAxis == 5) return g.gyro.z;

  return 0;
}

// Main Display Function
void updateDisplay(float value) {
  display.clearDisplay();

  // Sensor name and value
  if (state == 0) display.setCursor(0, 0), display.print("FLEX: ");
  if (state == 1) display.setCursor(0, 0), display.print("FORCE: ");
  if (state == 2) {
    const char* names[] = {"Acc X", "Acc Y", "Acc Z", "Gyro X", "Gyro Y", "Gyro Z"};
    display.setCursor(0, 0), display.print(names[mpuAxis]);
    display.print(": ");
  }
  display.print(value);

  display.display();
}

// Reads buttons with debounce
void readButtons() {
  if (digitalRead(BUTTON1_PIN) == LOW && millis() - lastDebounceTime1 > debounceDelay) {
    lastDebounceTime1 = millis();
    state = (state + 1) % 3;  // Cycle states (0-2)
    if (state == 2) mpuAxis = 0; // Reset MPU6050 axis on state change
    if (state == 0) Serial.println("Flex sensor:");
    if (state == 1) Serial.println("Force sensor:");
    if (state == 2) Serial.println("9DOF sensor:");
  }

  if (digitalRead(BUTTON2_PIN) == LOW && millis() - lastDebounceTime2 > debounceDelay) {
    lastDebounceTime2 = millis();
    if (state == 2) {
      mpuAxis = (mpuAxis + 1) % 6;  // Cycle MPU6050 axes
      if (mpuAxis == 0) Serial.println("X axis accelerometer");
      if (mpuAxis == 1) Serial.println("Y axis accelerometer");
      if (mpuAxis == 2) Serial.println("Z axis accelerometer");
      if (mpuAxis == 3) Serial.println("X axis gyroscope");
      if (mpuAxis == 4) Serial.println("Y axis gyroscope");
      if (mpuAxis == 5) Serial.println("Z axis gyroscope");
    }
  }
}

// Loop
void loop() {
  readButtons();
  float value = readSensor();
  updateDisplay(value);

  Serial.println(value);
  delay(50);
}