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

Adafruit_MPU6050 mpu;
Servo servoX; // Servo for X-axis
Servo servoY; // Servo for Y-axis

float angleX = 0;
float angleY = 0;
unsigned long lastTime;
float initialAccelAngleX = 0;
float initialAccelAngleY = 0;

void setup() {
  Serial.begin(115200);

  // Initialize MPU6050
  if (!mpu.begin()) {
    Serial.println("Failed to find MPU6050 chip");
    while (1) {
      delay(10); 
    }
  }
  mpu.setAccelerometerRange(MPU6050_RANGE_2_G);
  mpu.setGyroRange(MPU6050_RANGE_250_DEG);
  mpu.setFilterBandwidth(MPU6050_BAND_21_HZ);

  delay(100);

  // Capture the initial angles
  sensors_event_t a, g, temp;
  mpu.getEvent(&a, &g, &temp);
  initialAccelAngleX = atan2(a.acceleration.y, a.acceleration.z) * 180 / PI;
  initialAccelAngleY = atan2(-a.acceleration.x, sqrt(a.acceleration.y * a.acceleration.y + a.acceleration.z * a.acceleration.z)) * 180 / PI;

  lastTime = millis();

  // Initialize Servos
  servoX.attach(9); // Attach servo for X-axis to pin 9
  servoX.write(90); // Start at 90 degrees (neutral position)

  servoY.attach(10); // Attach servo for Y-axis to pin 10
  servoY.write(90); // Start at 90 degrees (neutral position)
}

void loop() {
  sensors_event_t a, g, temp;
  mpu.getEvent(&a, &g, &temp);

  // Calculate the elapsed time
  unsigned long currentTime = millis();
  float elapsedTime = (currentTime - lastTime) / 1000.0; // convert to seconds
  lastTime = currentTime;

  // Read accelerometer data
  float accelX = a.acceleration.x;
  float accelY = a.acceleration.y;
  float accelZ = a.acceleration.z;

  // Calculate the accelerometer angles
  float accelAngleX = atan2(accelY, accelZ) * 180 / PI - initialAccelAngleX;
  float accelAngleY = atan2(-accelX, sqrt(accelY * accelY + accelZ * accelZ)) * 180 / PI - initialAccelAngleY;

  // Read gyroscope data
  float gyroX = g.gyro.x; // angular velocity in rad/s
  float gyroY = g.gyro.y; // angular velocity in rad/s

  // Integrate the gyroscope data -> angle
  angleX += gyroX * elapsedTime;
  angleY += gyroY * elapsedTime;

  // Complementary filter to combine accelerometer and gyroscope data
  float alpha = 0.98;
  angleX = alpha * angleX + (1 - alpha) * accelAngleX;
  angleY = alpha * angleY + (1 - alpha) * accelAngleY;

  // Control the servos based on the angles to bring them back to 0 degrees
  int servoPosX = map(angleX, -45, 45, 0, 180);
  servoPosX = constrain(servoPosX, 0, 180);
  servoX.write(servoPosX);

  int servoPosY = map(angleY, -45, 45, 0, 180);
  servoPosY = constrain(servoPosY, 0, 180);
  servoY.write(servoPosY);

  // Print the angles and servo positions for debugging
  Serial.print("Angle X: ");
  Serial.print(angleX);
  Serial.print(" Angle Y: ");
  Serial.print(angleY);
  Serial.print(" Servo X Position: ");
  Serial.print(servoPosX);
  Serial.print(" Servo Y Position: ");
  Serial.println(servoPosY);

  delay(10); // Small delay for stability
}