#include <Wire.h>

// MPU6050 I2C address and registers
const int MPU_ADDR = 0x68;
const int PWR_MGMT_1 = 0x6B;
const int ACCEL_XOUT_H = 0x3B;
const int GYRO_XOUT_H = 0x43;

// LED pins
const int LED1 = 2;  // Compression depth reached (5 cm)
const int LED2 = 3;  // Return to top (0 cm)

// Variables for measurements
float accelX, accelY, accelZ;
float gyroX, gyroY, gyroZ;
float angleX = 0, angleY = 0;
float verticalAccel = 0;
float velocity = 0;
float displacement = 0;
unsigned long lastTime = 0;

// Thresholds
const float TARGET_DEPTH = 0.05;  // 5 cm in meters
const float RETURN_THRESHOLD = 0.01; // 1 cm tolerance for return-to-top
const float GRAVITY = 9.81;        // Earth's gravity in m/s²

// Complementary filter constants
const float ALPHA = 0.98;          // Gyro weight in complementary filter

void setup() {
  Serial.begin(115200);
  pinMode(LED1, OUTPUT);
  pinMode(LED2, OUTPUT);
  
  // Initialize MPU6050
  Wire.begin();
  Wire.beginTransmission(MPU_ADDR);
  Wire.write(PWR_MGMT_1);
  Wire.write(0); // Wake up the MPU6050
  Wire.endTransmission(true);
  
  Serial.println("CPR Depth Monitor with Tilt Compensation Started");
  Serial.println("Place sensor on chest and begin compressions");
}

void readMPU6050() {
  // Read accelerometer data
  Wire.beginTransmission(MPU_ADDR);
  Wire.write(ACCEL_XOUT_H);
  Wire.endTransmission(false);
  Wire.requestFrom(MPU_ADDR, 6, true);
  
  accelX = (Wire.read() << 8 | Wire.read()) / 16384.0;
  accelY = (Wire.read() << 8 | Wire.read()) / 16384.0;
  accelZ = (Wire.read() << 8 | Wire.read()) / 16384.0;
  
  // Read gyroscope data
  Wire.beginTransmission(MPU_ADDR);
  Wire.write(GYRO_XOUT_H);
  Wire.endTransmission(false);
  Wire.requestFrom(MPU_ADDR, 6, true);
  
  gyroX = (Wire.read() << 8 | Wire.read()) / 131.0;
  gyroY = (Wire.read() << 8 | Wire.read()) / 131.0;
  gyroZ = (Wire.read() << 8 | Wire.read()) / 131.0;
}

void loop() {
  // Read sensor data
  readMPU6050();
  
  // Calculate time difference
  unsigned long currentTime = micros();
  float deltaTime = (currentTime - lastTime) / 1000000.0; // Convert to seconds
  lastTime = currentTime;
  
  // Calculate angles using complementary filter
  // Accelerometer angle calculation
  float accelAngleX = atan2(accelY, accelZ) * RAD_TO_DEG;
  float accelAngleY = atan2(accelX, sqrt(accelY * accelY + accelZ * accelZ)) * RAD_TO_DEG;
  
  // Complementary filter to combine accelerometer and gyroscope data
  angleX = ALPHA * (angleX + gyroX * deltaTime) + (1 - ALPHA) * accelAngleX;
  angleY = ALPHA * (angleY + gyroY * deltaTime) + (1 - ALPHA) * accelAngleY;
  
  // Convert angles to radians for trigonometric functions
  float angleXRad = angleX * DEG_TO_RAD;
  float angleYRad = angleY * DEG_TO_RAD;
  
  // Calculate vertical acceleration using trigonometry
  // This removes the gravity component and compensates for tilt
  verticalAccel = accelZ * cos(angleXRad) * cos(angleYRad) - GRAVITY;
  
  // Integrate acceleration to get velocity
  velocity += verticalAccel * deltaTime;
  
  // Apply high-pass filter to velocity to reduce drift
  static float filteredVelocity = 0;
  filteredVelocity = 0.9 * filteredVelocity + 0.1 * velocity;
  velocity -= filteredVelocity * 0.1;
  
  // Integrate velocity to get displacement
  displacement += velocity * deltaTime;
  
  // Ensure displacement doesn't go negative
  if (displacement < 0) displacement = 0;
  
  // Check for target depth (5 cm)
  if (displacement >= TARGET_DEPTH) {
    digitalWrite(LED1, HIGH);
  } else {
    digitalWrite(LED1, LOW);
  }
  
  // Check for return to top
  if (displacement <= RETURN_THRESHOLD) {
    digitalWrite(LED2, HIGH);
    // Reset integration when at top to reduce drift
    velocity = 0;
  } else {
    digitalWrite(LED2, LOW);
  }
  
  // Display depth in centimeters
  Serial.print("Depth: ");
  Serial.print(displacement * 100); // Convert to cm
  Serial.print(" cm | Angle X: ");
  Serial.print(angleX);
  Serial.print("° | Angle Y: ");
  Serial.print(angleY);
  Serial.println("°");
  
  delay(50); // Short delay for stability
}