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

// MPU6050 and OLED setup
Adafruit_MPU6050 mpu;
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// Graph settings
#define GRAPH_WIDTH SCREEN_WIDTH
#define GRAPH_HEIGHT 32
#define GRAPH_OFFSET_Y 16

// Graph buffer
float graphBuffer[GRAPH_WIDTH];
int graphIndex = 0;

unsigned long lastMovementTime = 0;
float lastX = 0, lastY = 0, lastZ = 0;
bool isNoVibration = true;

void setup() {
  // Initialize Serial Monitor
  Serial.begin(115200);
  while (!Serial)
    delay(10);

  // Initialize MPU6050
  if (!mpu.begin()) {
    Serial.println("Failed to find MPU6050 chip");
    while (1) {
      delay(10);
    }
  }
  Serial.println("MPU6050 initialized!");

  // Set accelerometer range
  mpu.setAccelerometerRange(MPU6050_RANGE_2_G);

  // Initialize OLED (using I2C address 0x3D)
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3D)) { // Use 0x3D as the I2C address
    Serial.println("Failed to initialize OLED display");
    while (1);
  }
  display.clearDisplay();
  display.display();

  // Display startup message
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.println("Vibration Monitor");
  display.display();
  delay(2000);
  display.clearDisplay();
}

void loop() {
  // Get accelerometer event
  sensors_event_t accel, gyro, temp;
  mpu.getEvent(&accel, &gyro, &temp);

  // Read current accelerometer values
  float currentX = accel.acceleration.x;
  float currentY = accel.acceleration.y;
  float currentZ = accel.acceleration.z;

  // Check for movement
  if (fabs(currentX - lastX) > 0.01 || fabs(currentY - lastY) > 0.01 || fabs(currentZ - lastZ) > 0.01) {
    lastMovementTime = millis();
    isNoVibration = false;
  } else if (millis() - lastMovementTime > 50) {
    isNoVibration = true; // No vibration detected
  }

  // Update last positions
  lastX = currentX;
  lastY = currentY;
  lastZ = currentZ;

  // Calculate vibration magnitude in g (1 g = 9.81 m/s²)
  float vibrationMagnitude = isNoVibration ? 0 : sqrt(currentX * currentX + currentY * currentY + currentZ * currentZ) / 9.81;

  // Update graph buffer
  graphBuffer[graphIndex] = vibrationMagnitude;
  graphIndex = (graphIndex + 1) % GRAPH_WIDTH;

  // Draw graph and labels
  display.clearDisplay();
  drawGraph();
  drawLabels(vibrationMagnitude);
  display.display();

  // Send data to Serial Plotter
  Serial.print("Vibration: ");
  Serial.print(vibrationMagnitude, 2);
  Serial.println(" g");

  delay(50); // Adjust for smooth plotting
}

void drawGraph() {
  // Find max value in graphBuffer for scaling
  float maxVal = 0.01; // Small default to prevent divide by zero
  for (int i = 0; i < GRAPH_WIDTH; i++) {
    if (graphBuffer[i] > maxVal) {
      maxVal = graphBuffer[i];
    }
  }

  // Draw graph points
  for (int x = 0; x < GRAPH_WIDTH; x++) {
    // Calculate the display index relative to the graphIndex
    int index = (graphIndex + x) % GRAPH_WIDTH;

    // Map vibration magnitude to graph height
    int y = map(graphBuffer[index] * 100, 0, maxVal * 100, GRAPH_OFFSET_Y + GRAPH_HEIGHT, GRAPH_OFFSET_Y);

    // Draw the point
    display.drawPixel(x, y, SSD1306_WHITE);
  }

  // Draw axes
  display.drawLine(0, GRAPH_OFFSET_Y + GRAPH_HEIGHT, GRAPH_WIDTH, GRAPH_OFFSET_Y + GRAPH_HEIGHT, SSD1306_WHITE); // X-axis (Time)
  display.drawLine(0, GRAPH_OFFSET_Y, 0, GRAPH_OFFSET_Y + GRAPH_HEIGHT, SSD1306_WHITE); // Y-axis (Vibration)
}

void drawLabels(float vibrationMagnitude) {
  // Display Title
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.println("Vibration Monitor");

  // X-Axis Label (Time)
  display.setTextSize(1);
  display.setCursor(SCREEN_WIDTH - 50, GRAPH_OFFSET_Y + GRAPH_HEIGHT + 2);
  //display.println("Time");

  // Y-Axis Label (Vibration)
  display.setTextSize(1);
  display.setCursor(2, GRAPH_OFFSET_Y + GRAPH_HEIGHT - 8);
  //display.println("Vibration");

  // Display Current Vibration Value
  display.setCursor(0, SCREEN_HEIGHT - 10); // Position at the bottom
  display.print("Vibration: ");
  display.print(vibrationMagnitude, 2); // Show up to 2 decimal points
  display.println(" g");
}