#include <Adafruit_MPU6050.h>
#include <Adafruit_Sensor.h>
#include <Wire.h> // communication I2C
#include <Servo.h> // commande servos

Adafruit_MPU6050 mpu;

Servo servo_X;
Servo servo_Y;
Servo servo_Z;

int pos_X = 0;
int pos_Y = 0;
int pos_Z = 90;

int Gyro_out[3], Accel_out[3];

float dt = 0.02;

float Gyro_cal_x, Gyro_cal_y, Gyro_cal_z, Accel_cal_x, Accel_cal_y, Accel_cal_z;

float Gyro_X = 0;
float Accel_X = 0;
float Predicted_X = 0;

float Gyro_Y = 0;
float Accel_Y = 0;
float Predicted_Y = 0;

float Q = 2.5;
float R = 5;

float P00 = 0.1;
float P11 = 0.1;
float P01 = 0.1;

float Kk0, Kk1;

unsigned long timer;
unsigned long timex;
unsigned long lastMoveTime = 0; // Time of last movement

// PID control variables
float target_X = 0; // Desired position for X-axis
float target_Y = 0; // Desired position for Y-axis

float Kp = 100; // Proportional gain
float Ki = 1.2; // Integral gain
float Kd = 2; // Derivative gain

float integral_X = 0; // Integral term for X-axis
float integral_Y = 0; // Integral term for Y-axis

float prev_error_X = 0; // Previous error for X-axis
float prev_error_Y = 0; // Previous error for Y-axis

void setup()
{
  servo_X.attach(9);
  servo_Y.attach(10);
  // servo_Z.attach(10);

  int Gyro_cal_x_sample = 0;
  int Gyro_cal_y_sample = 0;
  int Gyro_cal_z_sample = 0;

  int Accel_cal_x_sample = 0;
  int Accel_cal_y_sample = 0;
  int Accel_cal_z_sample = 0;

  int i;

  delay(5);

  Wire.begin();
  Serial.begin(115200);

  mpu.begin();

  delay(100);

  for (i = 0; i < 100; i += 1)
  {
    sensors_event_t a, g, temp;
    mpu.getEvent(&a, &g, &temp);

    Accel_out[0] = a.acceleration.x;
    Accel_out[1] = a.acceleration.y;
    Accel_out[2] = a.acceleration.z;

    Gyro_out[0] = g.gyro.x;
    Gyro_out[1] = g.gyro.y;
    Gyro_out[2] = g.gyro.z;

    Gyro_cal_x_sample += Gyro_out[0];
    Gyro_cal_y_sample += Gyro_out[1];
    Gyro_cal_z_sample += Gyro_out[2];

    Accel_cal_x_sample += Accel_out[0];
    Accel_cal_y_sample += Accel_out[1];
    Accel_cal_z_sample += Accel_out[2];

    delay(50);
  }

  Gyro_cal_x = Gyro_cal_x_sample / 100;
  Gyro_cal_y = Gyro_cal_y_sample / 100;
  Gyro_cal_z = Gyro_cal_z_sample / 100;

  Accel_cal_x = Accel_cal_x_sample / 100;
  Accel_cal_y = Accel_cal_y_sample / 100;
  Accel_cal_z = (Accel_cal_z_sample / 100) - 256;
}

void loop()
{
  timer = millis();

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

  Accel_out[0] = a.acceleration.x;
  Accel_out[1] = a.acceleration.y;
  Accel_out[2] = a.acceleration.z;

  Gyro_out[0] = g.gyro.x;
  Gyro_out[1] = g.gyro.y;
  Gyro_out[2] = g.gyro.z;

  Accel_X = atan2((Accel_out[1] - Accel_cal_y) / 256, (Accel_out[2] - Accel_cal_z) / 256) * 180 / PI;
  Gyro_X = Gyro_X + ((Gyro_out[0] - Gyro_cal_x) / 14.375) * dt;

  if (Gyro_X < 180)
    Gyro_X += 360;
  if (Gyro_X >= 180)
    Gyro_X -= 360;

  Predicted_X = Predicted_X + ((Gyro_out[0] - Gyro_cal_x) / 14.375) * dt;

  Accel_Y = atan2((Accel_out[0] - Accel_cal_x) / 256, (Accel_out[2] - Accel_cal_z) / 256) * 180 / PI;
  Gyro_Y = Gyro_Y + ((Gyro_out[1] - Gyro_cal_y) / 14.375) * dt;

  if (Gyro_Y < 180)
    Gyro_Y += 360;
  if (Gyro_Y >= 180)
    Gyro_Y -= 360;

  Predicted_Y = Predicted_Y - ((Gyro_out[1] - Gyro_cal_y) / 14.375) * dt;

  P00 += dt * (2 * P01 + dt * P11);
  P01 += dt * P11;
  P00 += dt * Q;
  P11 += dt * Q;

  Kk0 = P00 / (P00 + R);
  Kk1 = P01 / (P01 + R);

  Predicted_X += (Accel_X - Predicted_X) * Kk0;
  Predicted_Y += (Accel_Y - Predicted_Y) * Kk0;

  P00 *= (1 - Kk0);
  P01 *= (1 - Kk1);
  P11 -= Kk1 * P01;

  // Reset integral term if no movement detected for a long time
  if ((millis() - lastMoveTime) > 5000) {
    integral_X = 0;
    integral_Y = 0;
  }

  // PID control calculations
  float error_X = target_X - Predicted_X;
  integral_X += error_X * dt;
  float derivative_X = (error_X - prev_error_X) / dt;
  float output_X = Kp * error_X + Ki * integral_X + Kd * derivative_X;

  float error_Y = target_Y - Predicted_Y;
  integral_Y += error_Y * dt;
  float derivative_Y = (error_Y - prev_error_Y) / dt;
  float output_Y = Kp * error_Y + Ki * integral_Y + Kd * derivative_Y;

  prev_error_X = error_X;
  prev_error_Y = error_Y;

  pos_X = map(output_X, -90, 90, 900, 2200);
  pos_Y = map(output_Y, 90, -90, 900, 2200);

  servo_X.writeMicroseconds(pos_X);
  servo_Y.writeMicroseconds(pos_Y);

  timex = millis() - timer;
  delay(20 - timex);
}