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

// Servo PWM Pins
#define SERVO_PITCH_PIN 9
#define SERVO_ROLL_PIN  10
#define SERVO_YAW_PIN   11

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
Adafruit_MPU6050 mpu;

// 3 Servo Motors
Servo servoPitch;
Servo servoRoll;
Servo servoYaw;

// Control Tuning Gains
const float K_PITCH       = 1.0;   // Static Pitch Angle compensation
const float K_ROLL        = 1.0;   // Static Roll Angle compensation
const float K_TREMOR_GYRO  = 10.0;  // Gyro tremor gain
const float K_TREMOR_ACCEL = 4.0;   // Accel tremor gain

// High-Pass Filter Variables
float emaAx = 0.0, emaAy = 0.0, emaAz = 9.8; 
float emaGx = 0.0, emaGy = 0.0, emaGz = 0.0;
const float alphaHP = 0.25; 

// 6-Axis Dynamic Tremor Outputs
float tremorAx = 0.0, tremorAy = 0.0, tremorAz = 0.0;
float tremorGx = 0.0, tremorGy = 0.0, tremorGz = 0.0;

// Calibration Offsets & Angles
float gyroXOffset = 0.0;
float gyroYOffset = 0.0;
float gyroZOffset = 0.0;
float pitchOffset = 0.0; 
float rollOffset  = 0.0; 

float pitch = 0.0, roll = 0.0;
const float alphaComp = 0.96;

int servoAnglePitch = 90;
int servoAngleRoll  = 90;
int servoAngleYaw   = 90;

bool oledActive = false;
unsigned long lastMicros = 0;

void setup() {
  Serial.begin(9600);
  delay(1000);

  Serial.println(F("\n=== 3-SERVO INDEPENDENT 6-AXIS STABILIZER STARTUP ==="));

  Wire.begin();
  Wire.setClock(100000);

  // 1. Initialize 3 Servos
  servoPitch.attach(SERVO_PITCH_PIN);
  servoRoll.attach(SERVO_ROLL_PIN);
  servoYaw.attach(SERVO_YAW_PIN);

  servoPitch.write(90);
  servoRoll.write(90);
  servoYaw.write(90);

  // 2. MPU6050 Init
  if (mpu.begin(0x68, &Wire)) {
    Serial.println(F("[MPU6050] Connected SUCCESS!"));
    mpu.setAccelerometerRange(MPU6050_RANGE_4_G);
    mpu.setGyroRange(MPU6050_RANGE_500_DEG);
    mpu.setFilterBandwidth(MPU6050_BAND_21_HZ);
  } else {
    Serial.println(F("[MPU6050] FAILED!"));
    while (1);
  }

  // 3. Sensor Baseline Auto-Calibration
  Serial.print(F("Calibrating Sensor Baseline... Keep Device FLAT & STILL! "));
  float gyroSumX = 0.0;
  float gyroSumY = 0.0;
  float gyroSumZ = 0.0;
  float pitchSum  = 0.0;
  float rollSum   = 0.0;

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

    float ax = a.acceleration.x;
    float ay = a.acceleration.y;
    float az = a.acceleration.z;

    gyroSumX += g.gyro.x;
    gyroSumY += g.gyro.y;
    gyroSumZ += g.gyro.z;

    pitchSum += atan2(-ax, sqrt(ay * ay + az * az)) * 180.0 / 3.14159265;
    rollSum  += atan2(ay, az) * 180.0 / 3.14159265;
    delay(10);
  }

  gyroXOffset = gyroSumX / 100.0;
  gyroYOffset = gyroSumY / 100.0;
  gyroZOffset = gyroSumZ / 100.0;
  pitchOffset = pitchSum / 100.0;
  rollOffset  = rollSum / 100.0;

  Serial.println(F("DONE!"));

  // 4. OLED Init
  if (display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    oledActive = true;
    display.clearDisplay();
    display.setTextColor(SSD1306_WHITE);
    display.setTextSize(1);
    display.setCursor(0, 0);
    display.println(F("SYSTEM ONLINE"));
    display.display();
  }

  lastMicros = micros();
}

void loop() {
  unsigned long nowMicros = micros();
  float dt = (nowMicros - lastMicros) / 1000000.0;
  if (dt <= 0.0 || dt > 0.1) dt = 0.01;
  lastMicros = nowMicros;

  sensors_event_t accel, gyro, temp;
  mpu.getEvent(&accel, &gyro, &temp);

  // Raw Accelerometer Values (m/s^2)
  float ax = accel.acceleration.x;
  float ay = accel.acceleration.y;
  float az = accel.acceleration.z;

  // Calibrated Gyroscope Values (deg/s) - ALL 3 AXES (X, Y, Z)
  float gxDeg = (gyro.gyro.x - gyroXOffset) * (180.0 / 3.14159265);
  float gyDeg = (gyro.gyro.y - gyroYOffset) * (180.0 / 3.14159265);
  float gzDeg = (gyro.gyro.z - gyroZOffset) * (180.0 / 3.14159265);

  // --- 1. ACCELEROMETER HIGH-PASS FILTER ---
  emaAx = (alphaHP * ax) + ((1.0 - alphaHP) * emaAx);
  emaAy = (alphaHP * ay) + ((1.0 - alphaHP) * emaAy);
  emaAz = (alphaHP * az) + ((1.0 - alphaHP) * emaAz);

  tremorAx = ax - emaAx;
  tremorAy = ay - emaAy;
  tremorAz = az - emaAz;

  // --- 2. GYROSCOPE HIGH-PASS FILTER ---
  emaGx = (alphaHP * gxDeg) + ((1.0 - alphaHP) * emaGx);
  emaGy = (alphaHP * gyDeg) + ((1.0 - alphaHP) * emaGy);
  emaGz = (alphaHP * gzDeg) + ((1.0 - alphaHP) * emaGz);

  tremorGx = gxDeg - emaGx;
  tremorGy = gyDeg - emaGy;
  tremorGz = gzDeg - emaGz;

  // --- 3. ORIENTATION FUSION (Pitch & Roll) ---
  float rawPitch = atan2(-ax, sqrt(ay * ay + az * az)) * 180.0 / 3.14159265 - pitchOffset;
  float rawRoll  = atan2(ay, az) * 180.0 / 3.14159265 - rollOffset;

  pitch = alphaComp * (pitch + gyDeg * dt) + (1.0 - alphaComp) * rawPitch;
  roll  = alphaComp * (roll + gxDeg * dt)  + (1.0 - alphaComp) * rawRoll;

  // --- 4. INDEPENDENT 3-SERVO CONTROL ---
  // Pitch Servo (Pin 9) -> Pitch Angle + Gyro Y Tremor
  float targetPitch = 90.0 - (K_PITCH * pitch) - (K_TREMOR_GYRO * tremorGy);
  servoAnglePitch = (int)constrain(targetPitch, 10, 170);
  servoPitch.write(servoAnglePitch);

  // Roll Servo (Pin 10) -> Roll Angle + Gyro X Tremor
  float targetRoll = 90.0 - (K_ROLL * roll) - (K_TREMOR_GYRO * tremorGx);
  servoAngleRoll = (int)constrain(targetRoll, 10, 170);
  servoRoll.write(servoAngleRoll);

  // Yaw Servo (Pin 11) -> Gyro Z Tremor
  float targetYaw = 90.0 - (K_TREMOR_GYRO * tremorGz);
  servoAngleYaw = (int)constrain(targetYaw, 10, 170);
  servoYaw.write(servoAngleYaw);

  // --- 5. FULL TELEMETRY STREAM WITH GYRO X, Y, Z INCLUDED ---
  Serial.print(F("Gyr[X,Y,Z]:"));
  Serial.print(gxDeg, 1); Serial.print(F(","));
  Serial.print(gyDeg, 1); Serial.print(F(","));
  Serial.print(gzDeg, 1);
  Serial.print(F(" |Acc[X,Y,Z]:"));
  Serial.print(tremorAx, 1); Serial.print(F(","));
  Serial.print(tremorAy, 1); Serial.print(F(","));
  Serial.print(tremorAz, 1);
  Serial.print(F(" | P:")); Serial.print((int)pitch);
  Serial.print(F(" R:")); Serial.print((int)roll);
  Serial.print(F(" | Servos[P,R,Y]:"));
  Serial.print(servoAnglePitch); Serial.print(F(","));
  Serial.print(servoAngleRoll); Serial.print(F(","));
  Serial.println(servoAngleYaw);

  // --- 6. OLED DISPLAY ---
  if (oledActive) {
    display.clearDisplay();
    
    // Line 1: Pitch & Roll angles
    display.setCursor(0, 0);
    display.print(F("P:")); display.print((int)pitch); display.print(F("deg"));
    display.setCursor(64, 0);
    display.print(F("R:")); display.print((int)roll); display.print(F("deg"));

    // Line 2: Accelerometer Tremors
    display.setCursor(0, 16);
    display.print(F("TrA:")); display.print(tremorAx, 1);
    display.setCursor(45, 16);
    display.print(F(",")); display.print(tremorAy, 1);
    display.setCursor(85, 16);
    display.print(F(",")); display.print(tremorAz, 1);

    // Line 3: Gyroscope Rates (Gx, Gy, Gz)
    display.setCursor(0, 32);
    display.print(F("Gyr:")); display.print(gxDeg, 1);
    display.setCursor(45, 32);
    display.print(F(",")); display.print(gyDeg, 1);
    display.setCursor(85, 32);
    display.print(F(",")); display.print(gzDeg, 1);

    // Line 4: Live 3-Servo Positions
    display.setCursor(0, 48);
    display.print(F("S[P,R,Y]:")); 
    display.print(servoAnglePitch); display.print(F(","));
    display.print(servoAngleRoll); display.print(F(","));
    display.print(servoAngleYaw);

    display.display();
  }

  delay(20);
}