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

//servo declarations
Servo servoX;
Servo servoY;

// Function declarations
void calibrate_gyro();
bool initializeMPU6050();
void gyro_signals();
void control_servo();
void pid_controller();
void serial_out();
void initialize_servo();

// PID Parameters
double Kp = 1.2;  // Proportional gain 1
double Ki = 0.1;  // Integral gain0.0.1
double Kd = 0.5;  // Derivative gain 0.5

// Setpoint Variables
double setpointX = 0.0;  // Desired angle for X-axis stabilization
double setpointY = 0.0;  // Desired angle for Y-axis stabilization

// Constants for Filtering
const float FILTER_ALPHA = 0.1;

// Constants for Integral Clamping
const double MAX_INTEGRAL = 1000.0;

// Constants for Serial Output Throttling
const unsigned long SERIAL_OUTPUT_INTERVAL = 1000;  // Time interval between serial outputs in milliseconds

// Variables for PID Control
double previousErrorX = 0.0, previousErrorY = 0.0;  // Previous error terms
double integralX = 0.0, integralY = 0.0;            // Integral terms
double pidOutputX, pidOutputY;
unsigned long lastTime = millis();
// Servo Positions
int servoPosX, servoPosY;

// Variables for Filtered PID Outputs
double filteredPidOutputX = 0.0;
double filteredPidOutputY = 0.0;

// Last Serial Output Time
unsigned long lastSerialOutputTime = 0;

// Sensor Data Variables
float RateRoll, RatePitch, RateYaw;
float AccX, AccY, AccZ;
float CalibrateAccX = -0.04, CalibrateAccY = 0.11, CalibrateAccZ = -0.04;
float AngleRoll, AnglePitch;

// Sensor Calibration Variables
float RateCalibrationRoll, RateCalibrationPitch, RateCalibrationYaw;
int RateCalibrationNumber;

// Define variables to store previous servo positions
int prevServoPosX = 90;  // Assuming the initial position is at 90
int prevServoPosY = 90;  // Assuming the initial position is at 90

void setup() {
  Serial.begin(9600);
  servoX.attach(9);   // Attach servo to pin 9
  servoY.attach(10);  // Attach servo to pin 10
  // calibrate_gyro();

  initialize_servo();
}

void loop() {
  gyro_signals();
  pid_controller();
  control_servo();
  unsigned long currentTime = millis();
  if (currentTime - lastSerialOutputTime >= SERIAL_OUTPUT_INTERVAL) {
    // Perform serial output
    serial_out();
    // Update last serial output time
    lastSerialOutputTime = currentTime;
  }
  delay(75);
}

void calibrate_gyro() {
  // Wire.setClock(400000);
  Wire.begin();
  delay(100);
  Wire.beginTransmission(0x68);
  Wire.write(0x6B);  // PWR_MGMT_1 register
  Wire.write(0);     // Wake up MPU6050
  Wire.endTransmission();
  for (RateCalibrationNumber = 0;
       RateCalibrationNumber < 500;
       RateCalibrationNumber++) {
    Serial.print("Calibration ");
    Serial.print(RateCalibrationNumber);
    Serial.println(" of 500");
    gyro_signals();
    RateCalibrationRoll += RateRoll;
    RateCalibrationPitch += RatePitch;
    RateCalibrationYaw += RateYaw;
    // delay(1);
  }
  RateCalibrationRoll /= 500;
  RateCalibrationPitch /= 500;
  RateCalibrationYaw /= 500;
}

bool initializeMPU6050() {
  Wire.begin();
  Wire.beginTransmission(0x68);
  Wire.write(0x6B);  // PWR_MGMT_1 register
  Wire.write(0);     // Wake up MPU6050
  if (Wire.endTransmission() != 0) {
    Serial.println("Error: Failed to initialize MPU6050.");

    return false;
  }

  // Calibration and other initialization tasks
  calibrate_gyro();

  return true;  // MPU6050 initialized successfully
}

void gyro_signals() {
  Wire.beginTransmission(0x68);
  Wire.write(0x3B);
  if (Wire.endTransmission() != 0) {
    Serial.println("Error: Failed to communicate with MPU6050.");
    return;
  }
  Wire.beginTransmission(0x68);
  Wire.write(0x1A);
  Wire.write(0x05);
  Wire.endTransmission();
  Wire.beginTransmission(0x68);
  Wire.write(0x1C);
  Wire.write(0x10);
  Wire.endTransmission();
  Wire.beginTransmission(0x68);
  Wire.write(0x3B);
  Wire.endTransmission();
  Wire.requestFrom(0x68, 6);
  int16_t AccXLSB = Wire.read() << 8 | Wire.read();
  int16_t AccYLSB = Wire.read() << 8 | Wire.read();
  int16_t AccZLSB = Wire.read() << 8 | Wire.read();
  Wire.beginTransmission(0x68);
  Wire.write(0x1B);
  Wire.write(0x8);
  Wire.endTransmission();
  Wire.beginTransmission(0x68);
  Wire.write(0x43);
  Wire.endTransmission();
  Wire.requestFrom(0x68, 6);
  int16_t GyroX = Wire.read() << 8 | Wire.read();
  int16_t GyroY = Wire.read() << 8 | Wire.read();
  int16_t GyroZ = Wire.read() << 8 | Wire.read();
  RateRoll = (float)GyroX / 65.5;
  RatePitch = (float)GyroY / 65.5;
  RateYaw = (float)GyroZ / 65.5;
  AccX = (float)AccXLSB / 4096 + CalibrateAccX;
  AccY = (float)AccYLSB / 4096 + CalibrateAccY;
  AccZ = (float)AccZLSB / 4096 + CalibrateAccZ;
  AngleRoll = atan(AccY / sqrt(AccX * AccX + AccZ * AccZ)) * 1 / (3.142 / 180);
  AnglePitch = -atan(AccX / sqrt(AccY * AccY + AccZ * AccZ)) * 1 / (3.142 / 180);
}

void control_servo() {
  // Apply low-pass filter to PID outputs
  filteredPidOutputX = (1 - FILTER_ALPHA) * filteredPidOutputX + FILTER_ALPHA * pidOutputX;
  filteredPidOutputY = (1 - FILTER_ALPHA) * filteredPidOutputY + FILTER_ALPHA * pidOutputY;

  // Map filtered PID outputs to servo positions
  int targetServoPosX = map(filteredPidOutputX, -90, 90, 0, 180);
  int targetServoPosY = map(filteredPidOutputY, -90, 90, 0, 180);

  // Ensure servo positions are within valid range
  targetServoPosX = constrain(targetServoPosX, 0, 180);
  targetServoPosY = constrain(targetServoPosY, 0, 180);

  servoPosX = prevServoPosX + filteredPidOutputX;
  servoPosY = prevServoPosY + filteredPidOutputY;

  servoPosX = constrain(servoPosX, 0, 180);
  servoPosY = constrain(servoPosY, 0, 180);

  // Update previous servo positions
  prevServoPosX = servoPosX;
  prevServoPosY = servoPosY;

  // Set servo positions
  servoX.write(servoPosX);
  servoY.write(servoPosY);
}

void pid_controller() {
  // Calculate errors with angle wrap-around handling
  double errorX = setpointX - AngleRoll;
  double errorY = setpointY - AnglePitch;
  // Handle angle wrap-around for errorX
  if (errorX > 180) {
    errorX -= 360;
  } else if (errorX < -180) {
    errorX += 360;
  }
  // Handle angle wrap-around for errorY
  if (errorY > 180) {
    errorY -= 360;
  } else if (errorY < -180) {
    errorY += 360;
  }

  // Update integral terms with clamping
  integralX += errorX;
  integralY += errorY;

  // Clamp integral terms
  integralX = constrain(integralX, -MAX_INTEGRAL, MAX_INTEGRAL);
  integralY = constrain(integralY, -MAX_INTEGRAL, MAX_INTEGRAL);

  // Calculate derivatives with respect to the sampling time
  unsigned long currentTime = millis();                  // Get time
  double deltaTime = (currentTime - lastTime) / 1000.0;  // Convert milliseconds to seconds
  double derivativeX = (errorX - previousErrorX) / deltaTime;
  double derivativeY = (errorY - previousErrorY) / deltaTime;

  // Update previous errors and last time
  previousErrorX = errorX;
  previousErrorY = errorY;
  lastTime = currentTime;

  // Compute PID outputs with filtered
  pidOutputX = Kp * errorX + Ki * integralX + Kd * derivativeX;
  pidOutputY = Kp * errorY + Ki * integralY + Kd * derivativeY;
}


void serial_out() {
  /* Print out the values */
  Serial.print("Acceleration X: ");
  Serial.print(AccX);
  Serial.print(", Y: ");
  Serial.print(AccY);
  Serial.print(", Z: ");
  Serial.print(AccZ);
  Serial.println(" g");

  Serial.print("Rotation X: ");
  Serial.print(RateRoll);
  Serial.print(", Y: ");
  Serial.print(RatePitch);
  Serial.print(", Z: ");
  Serial.print(RateYaw);
  Serial.println(" Degrees/s");

  Serial.print("Angle Roll: ");
  Serial.print(AngleRoll);
  Serial.print(", Pitch: ");
  Serial.print(AnglePitch);
  Serial.println(" Degrees");

  Serial.print("Error X: ");
  Serial.print(setpointX - AngleRoll);  // Calculate and print error for X-axis
  Serial.print(", Y: ");
  Serial.print(setpointY - AnglePitch);  // Calculate and print error for Y-axis
  Serial.println(" Degrees");

  Serial.print("Servo X: ");
  Serial.print(servoPosX);
  Serial.print(", Y: ");
  Serial.print(servoPosY);
  Serial.println(" Degrees");

  Serial.println("");
}

void initialize_servo() {
  servoX.write(0);
  delay(500);
  servoX.write(360);
  delay(500);
  servoX.write(90);
  delay(500);
  servoY.write(0);
  delay(500);
  servoY.write(360);
  delay(500);
  servoY.write(90);
  delay(500);
}