#include <Wire.h>
#include <Servo.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>

// ---------- OLED ----------
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_ADDR 0x3C
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// ---------- MPU6050 ----------
#define MPU_ADDR 0x68
#define REG_PWR_MGMT_1 0x6B
#define REG_ACCEL_XOUT_H 0x3B

// ---------- Servos ----------
#define SERVO_X_PIN 9   // X-axis tilt (left-right)  
#define SERVO_Y_PIN 10  // Y-axis tilt (forward-backward) 

Servo servoX;
Servo servoY;

// ---------- Timing (non-blocking) ----------
const unsigned long SENSOR_INTERVAL_MS = 20;   // ~50 Hz
const unsigned long OLED_INTERVAL_MS   = 100;  // 10 Hz screen update

unsigned long lastSensorMs = 0;
unsigned long lastOledMs   = 0;

// ---------- State ----------
float rollDeg = 0.0;   // X-axis inclination 
float pitchDeg = 0.0;  // Y-axis inclination

int servoXAngle = 90;
int servoYAngle = 90;
int lastServoXAngle = 90;
int lastServoYAngle = 90;

// ---------- Helpers ----------
static float clampf(float x, float lo, float hi) {
  if (x < lo) return lo;
  if (x > hi) return hi;
  return x;
}

static int clampi(int x, int lo, int hi) {
  if (x < lo) return lo;
  if (x > hi) return hi;
  return x;
}

// Map tilt degrees to servo degrees, with center at 90°.
// We lock tilt to ±45° for stable control.
int tiltToServo(float tiltDeg) {
  const float TILT_MAX = 45.0f;
  float t = clampf(tiltDeg, -TILT_MAX, TILT_MAX);
  // -45 -> 0, 0 -> 90, +45 -> 180
  float servo = (t + TILT_MAX) * (180.0f / (2.0f * TILT_MAX));
  return clampi((int)(servo + 0.5f), 0, 180);
}

const char* dirLeftRight(int angle) {
  if (angle > 95) return "RIGHT";
  if (angle < 85) return "LEFT";
  return "CENTER";
}

const char* dirForwardBack(int angle) {
  if (angle > 95) return "FORWARD";
  if (angle < 85) return "BACK";
  return "CENTER";
}

// ---------- MPU6050 low-level ----------
void mpuWrite(byte reg, byte value) {
  Wire.beginTransmission(MPU_ADDR);
  Wire.write(reg);
  Wire.write(value);
  Wire.endTransmission(true);
}

bool mpuReadAccelRaw(int16_t &ax, int16_t &ay, int16_t &az) {
  Wire.beginTransmission(MPU_ADDR);
  Wire.write(REG_ACCEL_XOUT_H);
  if (Wire.endTransmission(false) != 0) return false;

  // Request 6 bytes: ax, ay, az
  if (Wire.requestFrom(MPU_ADDR, (uint8_t)6, (uint8_t)true) != 6) return false;

  ax = (int16_t)((Wire.read() << 8) | Wire.read());
  ay = (int16_t)((Wire.read() << 8) | Wire.read());
  az = (int16_t)((Wire.read() << 8) | Wire.read());
  return true;
}

void mpuInit() {
  // Wake up MPU6050 (clear sleep bit)
  mpuWrite(REG_PWR_MGMT_1, 0x00);
}

// Compute roll/pitch from accelerometer only
// roll  = atan2(ay, az)
// pitch = atan2(-ax, sqrt(ay^2 + az^2))
void computeTiltFromAccel(int16_t axRaw, int16_t ayRaw, int16_t azRaw) {
  float ax = (float)axRaw;
  float ay = (float)ayRaw;
  float az = (float)azRaw;

  float roll  = atan2(ay, az);
  float pitch = atan2(-ax, sqrt(ay * ay + az * az));

  rollDeg  = roll  * 180.0f / PI;
  pitchDeg = pitch * 180.0f / PI;
}

// ---------- OLED ----------
void drawOLED() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);

  display.setCursor(0, 0);
  display.print("MPU6050 Tilt -> 2 Servos");

  display.setCursor(0, 16);
  display.print("Roll (X): ");
  display.print(rollDeg, 1);
  display.println(" deg");

  display.setCursor(0, 26);
  display.print("Servo X: ");
  display.print(servoXAngle);
  display.print("  ");
  display.println(dirLeftRight(servoXAngle));

  display.setCursor(0, 42);
  display.print("Pitch (Y): ");
  display.print(pitchDeg, 1);
  display.println(" deg");

  display.setCursor(0, 52);
  display.print("Servo Y: ");
  display.print(servoYAngle);
  display.print("  ");
  display.println(dirForwardBack(servoYAngle));

  display.display();
}

void setup() {
  Wire.begin();

  Serial.begin(115200);

  // Init MPU6050
  mpuInit();

  // Init servos
  servoX.attach(SERVO_X_PIN);
  servoY.attach(SERVO_Y_PIN);
  servoX.write(90);
  servoY.write(90);

  // Init OLED
  if (!display.begin(SSD1306_SWITCHCAPVCC, OLED_ADDR)) {
    // If OLED fails, keep running anyway
    Serial.println("OLED init failed");
  } else {
    display.clearDisplay();
    display.display();
  }
}

void loop() {
  unsigned long now = millis();

  // 1) Read sensor + update servo targets
  if (now - lastSensorMs >= SENSOR_INTERVAL_MS) {
    lastSensorMs = now;

    int16_t ax, ay, az;
    if (mpuReadAccelRaw(ax, ay, az)) {
      computeTiltFromAccel(ax, ay, az);

      // Independent control: X servo uses ONLY roll, Y servo uses ONLY pitch
      servoXAngle = tiltToServo(rollDeg);
      servoYAngle = tiltToServo(pitchDeg);

      // Reduce jitter: write only if meaningfully changed
      if (abs(servoXAngle - lastServoXAngle) >= 1) {
        servoX.write(servoXAngle);
        lastServoXAngle = servoXAngle;
      }
      if (abs(servoYAngle - lastServoYAngle) >= 1) {
        servoY.write(servoYAngle);
        lastServoYAngle = servoYAngle;
      }

      // Optional serial debug
      Serial.print("Roll=");
      Serial.print(rollDeg, 1);
      Serial.print(" Pitch=");
      Serial.print(pitchDeg, 1);
      Serial.print(" | SX=");
      Serial.print(servoXAngle);
      Serial.print(" SY=");
      Serial.println(servoYAngle);
    } else {
      Serial.println("MPU read failed");
    }
  }

  // 2) Update OLED periodically
  if (now - lastOledMs >= OLED_INTERVAL_MS) {
    lastOledMs = now;
    drawOLED();
  }
}