#include <AccelStepper.h>
#include <Wire.h> // Required for I2C communication
#include <MPU6050.h> // MPU6050 library

// Define Stepper Motor 1 pins (Left Wheel)
#define DIR_PIN_1 18
#define STEP_PIN_1 19

// Define Stepper Motor 2 pins (Right Wheel)
// --- CHANGED PINS FOR MPU6050 I2C ---
#define DIR_PIN_2 25
#define STEP_PIN_2 26

// Stepper object definition
AccelStepper stepper1(AccelStepper::DRIVER, STEP_PIN_1, DIR_PIN_1);
AccelStepper stepper2(AccelStepper::DRIVER, STEP_PIN_2, DIR_PIN_2);

// MPU6050 object
MPU6050 mpu;

// LED Pin (optional)
const int LED_PIN = 2;

// PID Constants (These will require tuning!)
// Start with these values and adjust them.
float Kp = 0.5; // Proportional gain
float Ki = 0.005; // Integral gain
float Kd = 0.05; // Derivative gain

// PID Variables
float setpoint = 0; // The desired "balance" angle (e.g., 0 degrees)
float error = 0;
float prevError = 0;
float integral = 0;
float derivative = 0;
float output = 0;

// MPU6050 sensor data
int16_t ax, ay, az;
int16_t gx, gy, gz;

// Angle estimation variables
float anglePitch = 0; // The angle we want to balance on (around the Y-axis)
float angleRoll = 0; // Angle around X-axis (sideways tilt)

// Complementary Filter parameters
float alpha = 0.98; // Weight for accelerometer (0.01 - 0.99). Higher means more accel trust.
                   // For self-balancing, usually high gyro trust (low accel trust) for short term
                   // and accel for long term. So, `alpha` (accel weight) should be small.
                   // Re-thinking: alpha is the gyro weight, 1-alpha is accel weight.
                   // Let's use it as Gyro_weight * (new_gyro_angle) + Accel_weight * (new_accel_angle)
                   // Often, it's Gyro_weight = 0.98, Accel_weight = 0.02
float dt_mpu = 0; // Delta time for MPU readings

// Loop timing for PID (different from MPU reading time)
unsigned long lastTimePID = 0;
float dt_pid = 0;

// Timing for MPU (to calculate dt_mpu)
unsigned long lastTimeMPU = 0;

void setup() {
  Serial.begin(115200);
  Serial.println("Self-Balancing Robot Simulation Starting with MPU6050...");

  // Initialize LED
  pinMode(LED_PIN, OUTPUT);
  digitalWrite(LED_PIN, LOW); // Start with LED off

  // Configure stepper motor properties
  stepper1.setMaxSpeed(500);
  stepper2.setMaxSpeed(500);
  stepper1.setAcceleration(200);
  stepper2.setAcceleration(200);
  stepper1.setSpeed(0);
  stepper2.setSpeed(0);

  // Initialize I2C
  Wire.begin();

  // Initialize MPU6050
  Serial.println("Initializing MPU6050...");
  mpu.initialize();
  // Verify connection
  Serial.print("MPU6050 connection: ");
  Serial.println(mpu.testConnection() ? "SUCCESS" : "FAILED");

  // Set gyro scale to 500 deg/s (or 2000 for faster movements)
  // MPU6050_GYRO_FS_250: 250 deg/s, 131 LSB/deg/s
  // MPU6050_GYRO_FS_500: 500 deg/s, 65.5 LSB/deg/s
  // MPU6050_GYRO_FS_1000: 1000 deg/s, 32.8 LSB/deg/s
  // MPU6050_GYRO_FS_2000: 2000 deg/s, 16.4 LSB/deg/s
  mpu.setFullScaleGyroRange(MPU6050_GYRO_FS_500);
  // Set accel scale to 2g (or 4g, 8g, 16g)
  // MPU6050_ACCEL_FS_2: 2g, 16384 LSB/g
  // MPU6050_ACCEL_FS_4: 4g, 8192 LSB/g
  // MPU6050_ACCEL_FS_8: 8g, 4096 LSB/g
  // MPU6050_ACCEL_FS_16: 16g, 2048 LSB/g
  mpu.setFullScaleAccelRange(MPU6050_ACCEL_FS_2);

  // Calibrate MPU (optional but recommended for real robots)
  // Wokwi simulation usually has perfect sensors, but good practice.
  // mpu.calibrateAccel();
  // mpu.calibrateGyro();

  // Take initial reading to seed complementary filter
  mpu.getMotion6(&ax, &ay, &az, &gx, &gy, &gz);
  // Calculate initial pitch angle from accelerometer (assuming y is forward, z is up/down)
  // Use atan2 for full 360 degree angle
  anglePitch = atan2(ax, az) * 180 / M_PI; // Pitch around Y-axis (tilt forward/backward)
  angleRoll = atan2(ay, az) * 180 / M_PI;   // Roll around X-axis (tilt side-to-side)

  Serial.print("Initial Pitch: ");
  Serial.println(anglePitch);

  lastTimePID = micros();
  lastTimeMPU = micros();
}

void loop() {
  // Calculate delta time for PID loop
  unsigned long currentTimePID = micros();
  dt_pid = (float)(currentTimePID - lastTimePID) / 1000000.0; // in seconds
  lastTimePID = currentTimePID;

  // Calculate delta time for MPU readings
  unsigned long currentTimeMPU = micros();
  dt_mpu = (float)(currentTimeMPU - lastTimeMPU) / 1000000.0; // in seconds
  lastTimeMPU = currentTimeMPU;

  // 1. Read MPU6050 data
  mpu.getMotion6(&ax, &ay, &az, &gx, &gy, &gz);

  // 2. Angle Estimation using Complementary Filter
  // Convert raw gyro data to degrees/second
  // For MPU6050_GYRO_FS_500, scale is 65.5 LSB/deg/s
  // Gyro X (roll rate) = gx / 65.5
  // Gyro Y (pitch rate) = gy / 65.5
  // Gyro Z (yaw rate) = gz / 65.5
  float gyroPitchRate = -(float)gy / 65.5; // Negative because positive GYRO Y is nose up
                                          // We want positive pitch when leaning forward

  // Calculate angle from accelerometer
  // For pitch (rotation around Y-axis), use AX and AZ
  // For roll (rotation around X-axis), use AY and AZ
  // The atan2(x,y) gives the angle relative to the positive x-axis.
  // For pitch: atan2(ax, az) gives angle relative to Z-axis.
  // If robot leans forward, AX becomes positive, AZ becomes less positive.
  float accelAnglePitch = atan2(ax, az) * 180 / M_PI;
  float accelAngleRoll = atan2(ay, az) * 180 / M_PI;

  // Complementary Filter for Pitch (the angle we care about for balancing)
  // anglePitch = (alpha * (anglePitch + gyroPitchRate * dt_mpu)) + ((1.0 - alpha) * accelAnglePitch);
  // Simpler approach (assuming alpha is gyro weight):
  anglePitch = alpha * (anglePitch + gyroPitchRate * dt_mpu) + (1.0 - alpha) * accelAnglePitch;


  Serial.print("Accel Pitch: ");
  Serial.print(accelAnglePitch);
  Serial.print(", Gyro Pitch Rate: ");
  Serial.print(gyroPitchRate);
  Serial.print(", Estimated Pitch: ");
  Serial.print(anglePitch);
  Serial.print(" deg");

  // 3. PID Calculation
  // Use the estimated pitch angle for the PID loop
  error = setpoint - anglePitch; // Error = Desired (0) - Current Angle

  integral += error * dt_pid;
  // Add anti-windup if needed for integral term
  // const float MAX_INTEGRAL = 100.0; // Adjust as needed
  // if (integral > MAX_INTEGRAL) integral = MAX_INTEGRAL;
  // if (integral < -MAX_INTEGRAL) integral = -MAX_INTEGRAL;

  // Use current gyro rate as derivative feedback, as it's cleaner than (error - prevError) / dt
  // This is common in balancing robots: P on angle, D on angular velocity
  derivative = -gyroPitchRate; // Negative because if robot is falling forward (positive gyroPitchRate),
                              // we want to counteract by moving motors forward, so output should be negative.
                              // (Error being positive when tilting back, output positive = forward)

  // Or, the standard (error - prevError) / dt_pid:
  // derivative = (error - prevError) / dt_pid;

  output = (Kp * error) + (Ki * integral) + (Kd * derivative);

  prevError = error; // Update prevError for next PID iteration

  Serial.print(", Error: ");
  Serial.print(error);
  Serial.print(", PID Output: ");
  Serial.println(output);

  // 4. Control Motors based on PID Output
  float motorSpeed = output * 20; // Scale the output to a reasonable motor speed range (adjust 20)
  motorSpeed = constrain(motorSpeed, -stepper1.maxSpeed(), stepper1.maxSpeed());

  // Set the speed for both motors
  // IMPORTANT: You might need to invert motorSpeed depending on your motor wiring and
  // how your robot should move to correct a tilt.
  // E.g., if leaning forward (positive anglePitch, negative error, negative output),
  // motors should move forward (positive speed). If this happens, then `motorSpeed` is correct.
  // If leaning back (negative anglePitch, positive error, positive output),
  // motors should move backward (negative speed).
  stepper1.setSpeed(motorSpeed);
  stepper2.setSpeed(motorSpeed);

  // Update motor positions (must be called continuously)
  stepper1.runSpeed();
  stepper2.runSpeed();

  // Simple LED feedback: turn on if the robot is "trying to balance"
  if (abs(motorSpeed) > 10) { // If motor is active
    digitalWrite(LED_PIN, HIGH);
  } else {
    digitalWrite(LED_PIN, LOW);
  }

  // Short delay for Wokwi and serial print stability
  delay(5);
}