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

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET    -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

#define FLEX_PIN      25    // Flex sensor analog input on GPIO25
#define LED_PIN       16    // LED output pin

#define MPU_UPPER_ADDR  0x68
#define MPU_LOWER_ADDR  0x69

Adafruit_MPU6050 mpu_upper;
Adafruit_MPU6050 mpu_lower;

unsigned long previousMillis = 0;
const long blinkInterval = 500; // Blink interval in ms
bool ledState = LOW;

const int flexSamples = 10;
float flexBuffer[flexSamples];
int flexIdx = 0;

void setup() {
  Serial.begin(115200);
  pinMode(LED_PIN, OUTPUT);
  digitalWrite(LED_PIN, LOW);

  Wire.begin();

  // Initialize OLED
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println("OLED display not found!");
    while (1);
  }
  display.clearDisplay();
  display.display();

  // Initialize MPUs
  if (!mpu_upper.begin(MPU_UPPER_ADDR)) {
    Serial.println("Upper MPU6050 not found!");
    while (1);
  }
  if (!mpu_lower.begin(MPU_LOWER_ADDR)) {
    Serial.println("Lower MPU6050 not found!");
    while (1);
  }

  // Fill flex buffer for startup
  for (int i = 0; i < flexSamples; i++) flexBuffer[i] = analogRead(FLEX_PIN);

  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0,0);
  display.println("Posture Monitor Init...");
  display.display();
  delay(1000);
}

// Get smoothed flex value (0–100)
float getSmoothFlex() {
  flexBuffer[flexIdx++] = analogRead(FLEX_PIN);
  if (flexIdx >= flexSamples) flexIdx = 0;
  float sum = 0;
  for (int i = 0; i < flexSamples; i++) sum += flexBuffer[i];
  float avg = sum / flexSamples;
  return avg * 100.0 / 4095.0;
}

void loop() {
  float flex_value = getSmoothFlex();

  sensors_event_t a_up, g_up, temp_up;
  mpu_upper.getEvent(&a_up, &g_up, &temp_up);

  sensors_event_t a_lo, g_lo, temp_lo;
  mpu_lower.getEvent(&a_lo, &g_lo, &temp_lo);

  // Calculate pitch angles (degrees, using x-axis)
  float pitch_up = atan2(-a_up.acceleration.x,
                         sqrt(pow(a_up.acceleration.y, 2) + pow(a_up.acceleration.z, 2))) * 180.0 / PI;
  float pitch_lo = atan2(-a_lo.acceleration.x,
                         sqrt(pow(a_lo.acceleration.y, 2) + pow(a_lo.acceleration.z, 2))) * 180.0 / PI;

  // Calibration: subtract initial sensor pitch at startup if needed
  // float pitch_offset_up = ...; float pitch_offset_lo = ...;
  // pitch_up -= pitch_offset_up;
  // pitch_lo -= pitch_offset_lo;

  // Thresholds: Only based on MPUs (x-axis/pitch)
  const float pitchThreshold = 20.0;
  bool bad_posture = abs(pitch_up) > pitchThreshold || abs(pitch_lo) > pitchThreshold;

  unsigned long currentMillis = millis();
  if (bad_posture) {
    if (currentMillis - previousMillis >= blinkInterval) {
      previousMillis = currentMillis;
      ledState = !ledState;
      digitalWrite(LED_PIN, ledState);
    }
  } else {
    digitalWrite(LED_PIN, LOW);
  }

  // OLED display
  display.clearDisplay();
  display.setTextSize(1);
  display.setCursor(0,0);
  display.print("Flex: "); display.println(flex_value, 1);
  display.print("Pitch Up: "); display.println(pitch_up, 1);
  display.print("Pitch Lo: "); display.println(pitch_lo, 1);
  display.setTextSize(2);
  display.setCursor(0,40);
  display.println(bad_posture ? "BAD POSTURE!" : "Posture OK");
  display.display();

  // Debug output
  Serial.print("Flex Value: "); Serial.print(flex_value, 1);
  Serial.print(" | Pitch Upper: "); Serial.print(pitch_up, 2);
  Serial.print(" | Pitch Lower: "); Serial.print(pitch_lo, 2);
  Serial.print(" | Status: "); Serial.println(bad_posture ? "BAD" : "OK");

  delay(50); // Loop timing
}