// CodeProven project R02 — supplied contributor source.
// Status: source supplied; bench verification and exact-board validation are required before production use.

// CodeProven project R02 — supplied contributor source.
// Status: source supplied; bench verification and exact-board validation are required before production use.

/*
  Project 2: Self-Balancing Rover (2-wheel inverted pendulum drive)
  Uses MPU6050 for tilt sensing + PID control driving two DC motors via L298N/TB6612.

  Hardware:
    - Arduino Uno/Nano
    - MPU6050 (I2C)
    - L298N motor driver
    - 2x DC gear motors + wheels
    - 7.4-12V battery for motors, separate from logic supply if possible

  Motor driver pins (L298N):
    ENA -> D5 (PWM), IN1 -> D6, IN2 -> D7   (Motor A)
    ENB -> D11(PWM), IN3 -> D8, IN4 -> D9   (Motor B)
*/

#include <Wire.h>
#include <Adafruit_MPU6050.h>
#include <Adafruit_Sensor.h>

Adafruit_MPU6050 mpu;

// Motor A
const int ENA = 5, IN1 = 6, IN2 = 7;
// Motor B
const int ENB = 11, IN3 = 8, IN4 = 9;

// PID constants - MUST be tuned per-build, these are a starting point
double Kp = 25.0;
double Ki = 140.0;
double Kd = 0.8;

double setpoint = 0.0;      // target angle (upright = 0, calibrate offset below)
double angle = 0.0;
double lastError = 0.0;
double integral = 0.0;
unsigned long lastTime = 0;

// Complementary filter
float filteredAngle = 0.0;
const float ALPHA = 0.98;

void setMotors(int speed) {
  // speed range: -255 to 255
  speed = constrain(speed, -255, 255);

  bool forward = speed >= 0;
  int pwm = abs(speed);

  digitalWrite(IN1, forward ? HIGH : LOW);
  digitalWrite(IN2, forward ? LOW  : HIGH);
  digitalWrite(IN3, forward ? HIGH : LOW);
  digitalWrite(IN4, forward ? LOW  : HIGH);

  analogWrite(ENA, pwm);
  analogWrite(ENB, pwm);
}

void setup() {
  Serial.begin(115200);
  pinMode(ENA, OUTPUT); pinMode(IN1, OUTPUT); pinMode(IN2, OUTPUT);
  pinMode(ENB, OUTPUT); pinMode(IN3, OUTPUT); pinMode(IN4, OUTPUT);

  Wire.begin();
  if (!mpu.begin()) {
    Serial.println("MPU6050 not found");
    while (1) delay(10);
  }
  mpu.setAccelerometerRange(MPU6050_RANGE_4_G);
  mpu.setGyroRange(MPU6050_RANGE_500_DEG);
  mpu.setFilterBandwidth(MPU6050_BAND_44_HZ);

  lastTime = millis();
  delay(1000);
}

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

  unsigned long now = millis();
  double dt = (now - lastTime) / 1000.0;
  if (dt <= 0) dt = 0.001;
  lastTime = now;

  // Angle from accelerometer (tilt around the axis the rover balances on)
  float accAngle = atan2(a.acceleration.x, a.acceleration.z) * 180.0 / PI;
  float gyroRate = g.gyro.y * 180.0 / PI;

  filteredAngle = ALPHA * (filteredAngle + gyroRate * dt) + (1 - ALPHA) * accAngle;
  angle = filteredAngle;

  // PID
  double error = setpoint - angle;
  integral += error * dt;
  integral = constrain(integral, -255, 255); // anti-windup
  double derivative = (error - lastError) / dt;
  double output = Kp * error + Ki * integral + Kd * derivative;
  lastError = error;

  // Safety cutoff: if it's fallen over too far, stop motors entirely
  if (abs(angle) > 45) {
    setMotors(0);
    integral = 0;
  } else {
    setMotors((int)output);
  }

  Serial.print("Angle: "); Serial.print(angle);
  Serial.print("  Output: "); Serial.println(output);

  delay(5); // ~200Hz loop - balancing needs a fast loop
}

/*
  CALIBRATION:
  1. Prop the rover perfectly vertical against a wall, read the printed
     "Angle" value, and set `setpoint` to that value (your true balance point,
     since IMU mounting is never perfectly aligned).
  2. Tune PID in this order: raise Kp until it oscillates gently in place,
     then add Kd to damp the oscillation, then add small Ki to kill steady
     drift. Retune every time you change weight distribution or wheel size.

  TESTING:
  - Start with wheels on the ground, rover propped at setpoint, and let go.
  - It should catch itself with small corrections, not violent motor slams.
  - If it oscillates wildly: reduce Kp. If it leans and slowly falls: increase Ki slightly.

  SAFETY:
  - Keep fingers clear of wheels during tuning - unexpected full-speed motor
    commands are common while tuning PID.
  - Use a fully charged battery; brownouts during tuning cause erratic behavior
    that looks like a control bug but isn't.
*/