#include "Arduino.h"
#include "Wire.h"
#include "MPU6050.h"
#include "AccelStepper.h"  // Library for controlling stepper motors
#include "NewPing.h"

#define TRIGGER_PIN 9
#define ECHO_PIN 8
#define MAX_DISTANCE 75

#define Kp  40
#define Kd  0.05
#define Ki  40
#define sampleTime  0.005
#define targetAngle -2.5

// Define Stepper Motor Pins
#define leftMotorStepPin 3
#define leftMotorDirPin  4
#define rightMotorStepPin 5
#define rightMotorDirPin 6

MPU6050 mpu;
int sampleCount = 1000;  // Number of samples to average
int16_t ax, ay, az, gx, gy, gz;
int32_t axSum = 0, aySum = 0, azSum = 0, gxSum = 0, gySum = 0, gzSum = 0;
int16_t axOffset, ayOffset, azOffset, gxOffset, gyOffset, gzOffset;

NewPing sonar(TRIGGER_PIN, ECHO_PIN, MAX_DISTANCE);

AccelStepper leftMotor(AccelStepper::DRIVER, leftMotorStepPin, leftMotorDirPin);
AccelStepper rightMotor(AccelStepper::DRIVER, rightMotorStepPin, rightMotorDirPin);

int16_t accY, accZ, gyroX;
volatile int motorPower, gyroRate;
volatile float accAngle, gyroAngle, currentAngle, prevAngle = 0, error, prevError = 0, errorSum = 0;
volatile byte count = 0;
int distanceCm;

void setup() {

   // Start serial communication for debugging
  Serial.begin(9600);
  // Initialize the MPU6050 sensor
  mpu.initialize();
  // Wait for sensor to stabilize (optional, but good practice)
  delay(1000);

 // Calibrate the MPU6050 (optional, if needed)
   calibrateMPU();
   
  // Apply the calculated offsets to the MPU6050
  mpu.setXGyroOffset(gxOffset);
  mpu.setYAccelOffset(ayOffset);
  mpu.setZAccelOffset(azOffset);

  // Print the offset values
  Serial.print("Gyro Offset X: ");
  Serial.println(gxOffset);
  Serial.print("Accel Offset Y: ");
  Serial.println(ayOffset);
  Serial.print("Accel Offset Z: ");
  Serial.println(azOffset);

  // Initialize the stepper motors
  leftMotor.setMaxSpeed(1000);  // Adjust speed as necessary
  leftMotor.setAcceleration(500);  // Adjust acceleration as necessary
  rightMotor.setMaxSpeed(1000);  // Adjust speed as necessary
  rightMotor.setAcceleration(500);  // Adjust acceleration as necessary

  // Set up the ultrasonic sensor
  sonar.ping_cm(); // Send a ping to clear the previous reading

  // Set motor control pins
  pinMode(13, OUTPUT);  // Status LED

  // Initialize PID sampling loop
  init_PID();
}

void loop() {
  // Read acceleration and gyroscope values
  accY = mpu.getAccelerationY();
  accZ = mpu.getAccelerationZ();
  gyroX = mpu.getRotationX();

  // Compute motor power based on PID control loop
  motorPower = constrain(motorPower, -255, 255);
  if (motorPower > 0) {
    leftMotor.setSpeed(motorPower);
    rightMotor.setSpeed(motorPower);
  } else {
    leftMotor.setSpeed(-motorPower);
    rightMotor.setSpeed(-motorPower);
  }

  // Measure distance every 100 milliseconds
  if ((count % 20) == 0) {
    distanceCm = sonar.ping_cm();
  }

  // If the robot is close to an obstacle, reverse direction
  if ((distanceCm < 20) && (distanceCm != 0)) {
    leftMotor.setSpeed(-motorPower);
    rightMotor.setSpeed(motorPower);
  }

  // Run the stepper motors
  leftMotor.run();
  Serial.println("Left");
  rightMotor.run();
  Serial.println("Right");
}

// Initialize Timer1 for PID control loop
void init_PID() {
  cli();          // Disable global interrupts
  TCCR1A = 0;     // Set entire TCCR1A register to 0
  TCCR1B = 0;     // Set entire TCCR1B register to 0
  OCR1A = 9999;   // Set compare match register to set sample time (5ms)
  TCCR1B |= (1 << WGM12); // Turn on CTC mode
  TCCR1B |= (1 << CS11);  // Set prescaler to 8
  TIMSK1 |= (1 << OCIE1A);  // Enable timer compare interrupt
  sei();          // Enable global interrupts
}

// Timer1 interrupt service routine for PID control (5ms loop)
ISR(TIMER1_COMPA_vect) {
  // Calculate the angle of inclination
  accAngle = atan2(accY, accZ) * RAD_TO_DEG;
  gyroRate = map(gyroX, -32768, 32767, -250, 250);
  gyroAngle = (float)gyroRate * sampleTime;
  currentAngle = 0.9934 * (prevAngle + gyroAngle) + 0.0066 * (accAngle);

  // Calculate error values for PID control
  error = currentAngle - targetAngle;
  errorSum = errorSum + error;
  errorSum = constrain(errorSum, -300, 300);

  // PID control to determine motor power
  motorPower = Kp * (error) + Ki * (errorSum) * sampleTime - Kd * (currentAngle - prevAngle) / sampleTime;
  prevAngle = currentAngle;

  // Toggle LED on pin13 every second (for debugging purposes)
  count++;
  if (count == 200) {
    count = 0;
    digitalWrite(13, !digitalRead(13));  // Toggle LED
  }
}

void calibrateMPU() {
  for (int i = 0; i < sampleCount; i++) {
    // Read accelerometer and gyroscope data
    mpu.getMotion6(&ax, &ay, &az, &gx, &gy, &gz);

    // Sum the accelerometer values
    axSum += ax;
    aySum += ay;
    azSum += az;

    // Sum the gyroscope values
    gxSum += gx;
    gySum += gy;
    gzSum += gz;

    // Small delay between readings to allow sensor to stabilize
    delay(10); 
  }

  // Calculate the offsets by averaging the values
  axOffset = axSum / sampleCount;
  ayOffset = aySum / sampleCount;
  azOffset = azSum / sampleCount;
  gxOffset = gxSum / sampleCount;
  gyOffset = gySum / sampleCount;
  gzOffset = gzSum / sampleCount;

  // Adjust for gravity on the Z-axis
  azOffset -= 16384; // Assuming the sensor is resting flat with Z-axis pointing up, adjust to zero gravity
  // For gyroscope offsets, no adjustment needed since the gyroscope should ideally read zero at rest.
}