#include <Wire.h>
#include <math.h>

// Delay for each main loop iteration in milliseconds
constexpr float kDelayTime = 5.0;

// Motor driver pin assignments
constexpr int kForwardPin = 8;   // Pin to drive motor forward
constexpr int kBackwardPin = 12; // Pin to drive motor backward
constexpr int kSpeedPin = 11;    // PWM pin for motor speed

// MPU6050 I2C address
constexpr int kMpuAddress = 0x68;

// State control flags
bool runGyroscope = true;    // Whether to read from the gyroscope
bool runMotor = false;       // Whether to run the motor
bool runDetector = false;    // Whether to track rotation
bool offsetsApplied = false; // Whether calibration is complete

// Orientation data
float pitch = 0, roll = 0;
float pitchOffset = 0, rollOffset = 0; // Accumulated offsets for calibration
float pitchZeroed = 0, rollZeroed = 0; // Zeroed pitch/roll values
float lastPitch = 0, lastRoll = 0;
float diffPitch = 0;                  // Difference in pitch for actuation
int gyroCounter = 0;                  // Counter for sampling pitch
int offsetCount = 0;                  // Counter for calibration samples
unsigned long startTime;             // Timestamp for calibration duration

// Switch input configuration
constexpr int kSwitchPin = 4;                     // Pin for toggle switch
constexpr unsigned long kDebounceDelay = 100;     // Debounce time
constexpr unsigned long kNextShiftDelay = 500;    // Minimum time between state shifts
bool systemActive = false;                        // System ON/OFF state
bool lastButtonState = HIGH;                      // Previous switch state
unsigned long lastDebounceTime = 0;
unsigned long lastShiftTime = 0;

// Rotation sensor input (e.g. hall-effect sensor)
constexpr int kRotationSensorPin = A0;
constexpr float kMotorRotationFactor = 3.0; // Empirical factor for translating counts to degrees
int motorRotationCount = 0;
int lastVoltage = 0;
bool startRotation = false;
unsigned long startRotationTime = 0;
int rotationCount = 0; // Counter for rotations

// Tilt sensor pins and states
constexpr int kTiltTopPin = 6;
constexpr int kTiltBotPin = 7;
bool tiltTopState = false;
bool tiltBotState = false;
bool lastTiltTopState = false;
bool lastTiltBotState = false;
bool boundTop = false;
bool boundBot = false;

// Motor direction flag (true = backward, false = forward)
bool motorDir = false;

// Function to control motor direction and activation
void controlMotor(float diffRoll, bool run, bool& direction, bool& detector) {
  if (run) {
    // Decide direction based on sign of pitch change
    if (diffRoll > 0) {
      digitalWrite(kForwardPin, LOW);
      digitalWrite(kBackwardPin, HIGH);
      direction = false;
    } else {
      digitalWrite(kForwardPin, HIGH);
      digitalWrite(kBackwardPin, LOW);
      direction = true;
    }
    analogWrite(kSpeedPin, 255); // Run motor at full speed
    detector = true;             // Enable rotation detection
  } else {
    digitalWrite(kForwardPin, LOW);
    digitalWrite(kBackwardPin, LOW); // Stop motor
  }
}

// Function to detect motor rotation using a hall-effect sensor
void detectRotation() {
  int currentVoltage = analogRead(kRotationSensorPin) * (5.0 / 1023.0);

  if (!startRotation) {
    startRotationTime = millis();
    startRotation = true; // Begin tracking duration
  }

  // Detect rising edge on hall sensor signal
  if (currentVoltage != 0 && lastVoltage == 0) {
    rotationCount++;
    Serial.print("Current rotation: ");
    Serial.println(rotationCount);

    float rotationThreshold = abs(diffPitch);

    // Stop motor if rotation target reached
    if (rotationCount * kMotorRotationFactor> rotationThreshold) {
      runMotor = false;
      runGyroscope = true;
      runDetector = false;
      startRotation = false;
      rotationCount = 0;
      controlMotor(0, false, motorDir, runDetector);
    }
  }

  // Timeout failsafe to prevent endless motor run
  if (millis() - startRotationTime > 4000) {
    runMotor = false;
    runGyroscope = true;
    runDetector = false;
    startRotation = false;
    rotationCount = 0;
    controlMotor(0, false, motorDir, runDetector);
    Serial.println("Rotation timeout");
  }

  lastVoltage = currentVoltage;
  checkTiltBoundary();

  // Check boundaries again for safety
  if (boundTop || boundBot) {
    runMotor = false;
    runGyroscope = true;
    runDetector = false;
    startRotation = false;
    rotationCount = 0;
    controlMotor(0, false, motorDir, runDetector);
  }
}

// Checks if either tilt sensor has reached boundary
void checkTiltBoundary() {
  tiltTopState = digitalRead(kTiltTopPin);
  tiltBotState = digitalRead(kTiltBotPin);

  // Trigger top limit if moving downward
  if (!lastTiltTopState && !lastTiltBotState && tiltTopState && !tiltBotState && diffPitch < 0) {
    boundTop = true;
    Serial.println("Top boundary reached");
  }
  // Trigger bottom limit if moving upward
  else if (!lastTiltTopState && lastTiltBotState && !tiltTopState && !tiltBotState && diffPitch > 0) {
    boundBot = true;
    Serial.println("Bottom boundary reached");
  } else {
    boundTop = false;
    boundBot = false;
  }

  lastTiltTopState = tiltTopState;
  lastTiltBotState = tiltBotState;
}

// Reads raw accelerometer data from MPU6050 and calculates pitch and roll
void readMpu6050(int outputEvery) {
  Wire.beginTransmission(kMpuAddress);
  Wire.write(0x3B); // Starting register for accelerometer
  Wire.endTransmission(false);
  Wire.requestFrom(kMpuAddress, 6, true);

  // Read accelerometer data (X, Y, Z)
  int16_t accX = Wire.read() << 8 | Wire.read();
  int16_t accY = Wire.read() << 8 | Wire.read();
  int16_t accZ = Wire.read() << 8 | Wire.read();

  // Convert to G-units
  float ax = accX / 16384.0;
  float ay = accY / 16384.0;
  float az = accZ / 16384.0;

  // Compute pitch and roll from acceleration
  pitch = atan2(ay, sqrt(ax * ax + az * az)) * 180.0 / PI;
  roll = atan2(-ax, az) * 180.0 / PI;

  // Perform calibration for first 3 seconds
  if (millis() - startTime < 3000) {
    pitchOffset += pitch;
    rollOffset += roll;
    offsetCount++;
    if (offsetCount % 100 == 0) Serial.println("Calibrating...");
  } else if (!offsetsApplied && offsetCount > 0) {
    pitchOffset /= offsetCount;
    rollOffset /= offsetCount;
    offsetsApplied = true;
    Serial.println("Calibration complete");
  }

  // Use calibrated data after offset applied
  if (offsetsApplied) {
    pitchZeroed = pitch - pitchOffset;
    rollZeroed = roll - rollOffset;

    if (gyroCounter == outputEvery) {
      diffPitch = lastPitch - pitchZeroed;
      Serial.print("Pitch Diff: ");
      Serial.println(diffPitch);

      if (abs(diffPitch) > 1.5) {
        runMotor = true;
        runGyroscope = false;
      }

      Serial.print("Pitch: ");
      Serial.println(pitchZeroed);
      lastPitch = pitchZeroed;
      gyroCounter = 0;
    }
  }

  gyroCounter++;
}

// Toggles the system state using a debounced switch
void handleSwitchToggle() {
  bool reading = digitalRead(kSwitchPin);

  if (reading != lastButtonState) {
    lastDebounceTime = millis();
  }

  if ((millis() - lastDebounceTime > kDebounceDelay) && (millis() - lastShiftTime > kNextShiftDelay)) {
    if (reading == LOW) {
      systemActive = !systemActive;
      lastShiftTime = millis();
      Serial.print("System is now: ");
      Serial.println(systemActive ? "ON" : "OFF");
      delay(500); // Prevent bouncing
    }
  }

  lastButtonState = reading;
}

// Setup function to initialize sensors, pins, and I2C
void setup() {
  pinMode(kForwardPin, OUTPUT);
  pinMode(kBackwardPin, OUTPUT);
  pinMode(kSpeedPin, OUTPUT);
  pinMode(kSwitchPin, INPUT_PULLUP);
  pinMode(kTiltTopPin, INPUT);
  pinMode(kTiltBotPin, INPUT);

  Serial.begin(9600);
  Wire.begin();
  Wire.beginTransmission(kMpuAddress);
  Wire.write(0x6B); // Power management register
  Wire.write(0);    // Wake up MPU6050
  Wire.endTransmission(true);

  startTime = millis();
}

// Main loop to coordinate gyro reading, motor control, and sensor input
void loop() {
  handleSwitchToggle();

  if (runGyroscope) {
    readMpu6050(50); // Check orientation every 50 cycles
  } else if (runMotor && systemActive) {
    if (runDetector) {
      detectRotation();
    } else {
      controlMotor(diffPitch, runMotor, motorDir, runDetector);
    }
  }

  delay(kDelayTime);
}