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

// =====================================================
// PIN DEFINITIONS
// =====================================================

const int MPU_ADDR = 0x68;

const int BUTTON_PIN = 2;

const int RED_PIN = 9;
const int GREEN_PIN = 10;
const int BLUE_PIN = 11;


// =====================================================
// CALIBRATION SETTINGS
// =====================================================

const int CALIBRATION_REPS = 5;

// How many standard deviations are included
// in the personalized tolerance region.
const float TOLERANCE_SD = 2.0;

// Minimum pitch excursion required for a valid rep.
const float MIN_REP_RANGE = 15.0;

// Minimum pitch change needed to recognize
// that the user has left the bottom position.
const float MOTION_THRESHOLD = 5.0;

// Minimum amount the pitch must reverse after
// reaching the top before recognizing the top.
const float TOP_REVERSAL_THRESHOLD = 3.0;

// Time the user must remain still at the bottom
// when calibration begins.
const unsigned long BOTTOM_SETUP_TIME = 1500;


// =====================================================
// BUTTON SETTINGS
// =====================================================

const unsigned long DEBOUNCE_DELAY = 50;

// Holding the button this long starts calibration.
const unsigned long LONG_PRESS_TIME = 2000;


// =====================================================
// ORIENTATION
// =====================================================

const float PITCH_SIGN = 1.0;
const float ROLL_SIGN = 1.0;


// =====================================================
// SYSTEM STATES
// =====================================================

bool systemOn = false;
bool calibrationMode = false;
bool calibrationComplete = false;


// =====================================================
// BUTTON VARIABLES
// =====================================================

bool lastButtonReading = HIGH;
bool buttonState = HIGH;

unsigned long lastDebounceTime = 0;
unsigned long buttonPressStart = 0;

bool longPressHandled = false;


// =====================================================
// GYROSCOPE VARIABLES
// =====================================================

float gyroOffsetX = 0;
float gyroOffsetY = 0;
float gyroOffsetZ = 0;

float pitch = 0;
float roll = 0;

float previousTime = 0;


// =====================================================
// REP STATE
// =====================================================

enum RepState {
  AT_BOTTOM,
  MOVING_UP,
  AT_TOP,
  MOVING_DOWN
};

RepState repState = AT_BOTTOM;


// =====================================================
// CALIBRATION REP DATA
// =====================================================

int calibrationRepCount = 0;

float repMinPitch = 0;
float repMaxPitch = 0;

float bottomSum = 0;
float topSum = 0;

float bottomSumSq = 0;
float topSumSq = 0;


// =====================================================
// PERSONALIZED BASELINE
// =====================================================

float baselineBottom = 0;
float baselineTop = 0;

float bottomSD = 0;
float topSD = 0;

float bottomTolerance = 5;
float topTolerance = 5;


// =====================================================
// MOVEMENT VARIABILITY
// =====================================================

float previousGyroX = 0;
float previousGyroY = 0;
float previousGyroZ = 0;

float liveShake = 0;

float shakeSum = 0;
float shakeSumSq = 0;

float baselineShake = 0;
float baselineShakeSD = 1;


// =====================================================
// LED
// =====================================================

void setLED(int red, int green, int blue) {

  analogWrite(RED_PIN, red);
  analogWrite(GREEN_PIN, green);
  analogWrite(BLUE_PIN, blue);
}


// =====================================================
// MPU6050 SETUP
// =====================================================

void setupMPU() {

  // Wake up MPU6050
  Wire.beginTransmission(MPU_ADDR);
  Wire.write(0x6B);
  Wire.write(0);
  Wire.endTransmission(true);

  // Accelerometer: ±2g
  Wire.beginTransmission(MPU_ADDR);
  Wire.write(0x1C);
  Wire.write(0x00);
  Wire.endTransmission(true);

  // Gyroscope: ±250 deg/s
  Wire.beginTransmission(MPU_ADDR);
  Wire.write(0x1B);
  Wire.write(0x00);
  Wire.endTransmission(true);
}


// =====================================================
// READ MPU6050
// =====================================================

void readMPU(
  int16_t &ax,
  int16_t &ay,
  int16_t &az,
  int16_t &gx,
  int16_t &gy,
  int16_t &gz
) {

  Wire.beginTransmission(MPU_ADDR);
  Wire.write(0x3B);
  Wire.endTransmission(false);

  Wire.requestFrom(MPU_ADDR, 14, true);

  ax = Wire.read() << 8 | Wire.read();
  ay = Wire.read() << 8 | Wire.read();
  az = Wire.read() << 8 | Wire.read();

  // Temperature
  Wire.read();
  Wire.read();

  gx = Wire.read() << 8 | Wire.read();
  gy = Wire.read() << 8 | Wire.read();
  gz = Wire.read() << 8 | Wire.read();
}


// =====================================================
// GYROSCOPE CALIBRATION
// =====================================================

void calibrateGyro() {

  long sumX = 0;
  long sumY = 0;
  long sumZ = 0;

  const int samples = 500;

  Serial.println("Keep sensor completely still...");
  
  for (int i = 0; i < samples; i++) {

    int16_t ax, ay, az;
    int16_t gx, gy, gz;

    readMPU(ax, ay, az, gx, gy, gz);

    sumX += gx;
    sumY += gy;
    sumZ += gz;

    delay(2);
  }

  gyroOffsetX = sumX / (float)samples;
  gyroOffsetY = sumY / (float)samples;
  gyroOffsetZ = sumZ / (float)samples;

  Serial.println("Gyroscope calibrated.");
}


// =====================================================
// ACCELEROMETER PITCH
// =====================================================

float getAccelerometerPitch(
  int16_t ax,
  int16_t ay,
  int16_t az
) {

  float ax_g = ax / 16384.0;
  float ay_g = ay / 16384.0;
  float az_g = az / 16384.0;

  float accelPitch =
    atan2(
      ax_g,
      sqrt(ay_g * ay_g + az_g * az_g)
    )
    * 180.0 / PI;

  return accelPitch * PITCH_SIGN;
}


// =====================================================
// RESET CALIBRATION
// =====================================================

void resetCalibration() {

  calibrationMode = true;
  calibrationComplete = false;

  calibrationRepCount = 0;

  bottomSum = 0;
  topSum = 0;

  bottomSumSq = 0;
  topSumSq = 0;

  baselineBottom = 0;
  baselineTop = 0;

  bottomSD = 0;
  topSD = 0;

  bottomTolerance = 5;
  topTolerance = 5;

  shakeSum = 0;
  shakeSumSq = 0;

  baselineShake = 0;
  baselineShakeSD = 1;

  repState = AT_BOTTOM;

  liveShake = 0;

  setLED(0, 0, 255);

  Serial.println();
  Serial.println("=================================");
  Serial.println("CALIBRATION MODE");
  Serial.println("Hold the weight at the BOTTOM.");
  Serial.println("Starting calibration...");
  Serial.println("=================================");
}


// =====================================================
// CHECK BOTTOM REGION
// =====================================================

bool isAtBottom() {

  return abs(pitch - baselineBottom)
         <= bottomTolerance;
}


// =====================================================
// CHECK TOP REGION
// =====================================================

bool isAtTop() {

  return abs(pitch - baselineTop)
         <= topTolerance;
}


// =====================================================
// RECORD CALIBRATION REP
// =====================================================

void recordCalibrationRep() {

  float repBottom = repMinPitch;
  float repTop = repMaxPitch;

  float range = repTop - repBottom;

  // Reject tiny movements
  if (range < MIN_REP_RANGE) {

    Serial.println("Movement too small. Rep ignored.");
    return;
  }

  calibrationRepCount++;

  bottomSum += repBottom;
  topSum += repTop;

  bottomSumSq += repBottom * repBottom;
  topSumSq += repTop * repTop;

  shakeSum += liveShake;
  shakeSumSq += liveShake * liveShake;

  Serial.print("Calibration Rep ");
  Serial.print(calibrationRepCount);

  Serial.print(" | Bottom = ");
  Serial.print(repBottom);

  Serial.print(" | Top = ");
  Serial.println(repTop);

  // Reset tracking for next rep
  repMinPitch = pitch;
  repMaxPitch = pitch;

  repState = AT_BOTTOM;

  // ===================================================
  // FINISHED 10 REPS
  // ===================================================

  if (calibrationRepCount >= CALIBRATION_REPS) {

    // Calculate averages
    baselineBottom =
      bottomSum / CALIBRATION_REPS;

    baselineTop =
      topSum / CALIBRATION_REPS;

    baselineShake =
      shakeSum / CALIBRATION_REPS;


    // Calculate variance
    float bottomVariance =
      (bottomSumSq / CALIBRATION_REPS)
      - (baselineBottom * baselineBottom);

    float topVariance =
      (topSumSq / CALIBRATION_REPS)
      - (baselineTop * baselineTop);

    float shakeVariance =
      (shakeSumSq / CALIBRATION_REPS)
      - (baselineShake * baselineShake);


    // Calculate standard deviations
    bottomSD =
      sqrt(max(0.0, bottomVariance));

    topSD =
      sqrt(max(0.0, topVariance));

    baselineShakeSD =
      sqrt(max(0.0, shakeVariance));


    // Prevent extremely small tolerances
    if (bottomSD < 1.5) {
      bottomSD = 1.5;
    }

    if (topSD < 1.5) {
      topSD = 1.5;
    }

    if (baselineShakeSD < 0.1) {
      baselineShakeSD = 0.1;
    }


    // Create personalized regions
    bottomTolerance =
      bottomSD * TOLERANCE_SD;

    topTolerance =
      topSD * TOLERANCE_SD;


    // Minimum tolerance
    if (bottomTolerance < 3.0) {
      bottomTolerance = 3.0;
    }

    if (topTolerance < 3.0) {
      topTolerance = 3.0;
    }


    calibrationComplete = true;
    calibrationMode = false;

    repState = AT_BOTTOM;

    setLED(0, 255, 0);


    Serial.println();
    Serial.println("=================================");
    Serial.println("CALIBRATION COMPLETE");
    Serial.println("=================================");

    Serial.print("Bottom = ");
    Serial.print(baselineBottom);

    Serial.print(" +/- ");
    Serial.println(bottomTolerance);

    Serial.print("Top = ");
    Serial.print(baselineTop);

    Serial.print(" +/- ");
    Serial.println(topTolerance);

    Serial.print("Movement variability = ");
    Serial.println(baselineShake);

    Serial.println();
    Serial.println("SYSTEM READY");
    Serial.println();
  }
}


// =====================================================
// CALIBRATION REP DETECTION
// =====================================================

void detectCalibrationRep() {

  // Continuously track the extremes
  if (pitch < repMinPitch) {
    repMinPitch = pitch;
  }

  if (pitch > repMaxPitch) {
    repMaxPitch = pitch;
  }


  // ---------------------------------------------
  // USER STARTS MOVING
  // ---------------------------------------------

  if (repState == AT_BOTTOM) {

    if (abs(pitch - repMinPitch)
        > MOTION_THRESHOLD) {

      repState = MOVING_UP;
    }
  }


  // ---------------------------------------------
  // USER REACHES TOP AND STARTS REVERSING
  // ---------------------------------------------

  else if (repState == MOVING_UP) {

    if (pitch <
        repMaxPitch - TOP_REVERSAL_THRESHOLD) {

      repState = AT_TOP;
    }
  }


  // ---------------------------------------------
  // USER STARTS MOVING DOWN
  // ---------------------------------------------

  else if (repState == AT_TOP) {

    if (pitch <
        repMaxPitch - TOP_REVERSAL_THRESHOLD) {

      repState = MOVING_DOWN;
    }
  }


  // ---------------------------------------------
  // USER RETURNS TO BOTTOM
  // ---------------------------------------------

  else if (repState == MOVING_DOWN) {

    if (abs(pitch - repMinPitch)
        < MOTION_THRESHOLD) {

      recordCalibrationRep();
    }
  }
}


// =====================================================
// NORMAL REP DETECTION
// =====================================================

void detectNormalRep() {

  switch (repState) {

    // ---------------------------------------------
    // WAITING AT BOTTOM
    // ---------------------------------------------

    case AT_BOTTOM:

      if (!isAtBottom()) {

        repState = MOVING_UP;

        Serial.println("Moving up");
      }

      break;


    // ---------------------------------------------
    // MOVING UP
    // ---------------------------------------------

    case MOVING_UP:

      if (isAtTop()) {

        repState = AT_TOP;

        Serial.println("Top reached");
      }

      break;


    // ---------------------------------------------
    // AT TOP
    // ---------------------------------------------

    case AT_TOP:

      if (!isAtTop()) {

        repState = MOVING_DOWN;

        Serial.println("Moving down");
      }

      break;


    // ---------------------------------------------
    // MOVING DOWN
    // ---------------------------------------------

    case MOVING_DOWN:

      if (isAtBottom()) {

        Serial.println("REP COMPLETE");

        repState = AT_BOTTOM;
      }

      break;
  }
}


// =====================================================
// MOVEMENT FEEDBACK
// =====================================================

void updateFeedback() {

  // Calibration = blue
  if (calibrationMode) {

    setLED(0, 0, 255);

    return;
  }


  // Not calibrated
  if (!calibrationComplete) {

    setLED(0, 0, 0);

    return;
  }


  // Compare current movement variability
  // to the user's baseline.

  float deviation =
    liveShake - baselineShake;


  // ---------------------------------------------
  // NORMAL
  // ---------------------------------------------

  if (deviation <= 0) {

    setLED(0, 255, 0);
  }


  // ---------------------------------------------
  // INCREASING DEVIATION
  // ---------------------------------------------

  else {

    float zScore =
      deviation / baselineShakeSD;

    float intensity =
      constrain(zScore / 3.0, 0, 1);


    // Green → Yellow
    if (intensity < 0.5) {

      int red =
        intensity * 2 * 255;

      setLED(red, 255, 0);
    }


    // Yellow → Red
    else {

      int green =
        (1 - intensity) * 2 * 255;

      setLED(255, green, 0);
    }
  }
}


// =====================================================
// BUTTON HANDLING
// =====================================================

void handleButton() {

  bool reading =
    digitalRead(BUTTON_PIN);


  // Detect electrical state change
  if (reading != lastButtonReading) {

    lastDebounceTime = millis();
  }


  // Wait for debounce
  if ((millis() - lastDebounceTime)
      > DEBOUNCE_DELAY) {

    if (reading != buttonState) {

      buttonState = reading;


      // =========================================
      // BUTTON PRESSED
      // =========================================

      if (buttonState == LOW) {

        buttonPressStart = millis();

        longPressHandled = false;
      }


      // =========================================
      // BUTTON RELEASED
      // =========================================

      else {

        unsigned long pressDuration =
          millis() - buttonPressStart;


        // If it wasn't already handled as
        // a long press, treat it as a short press.

        if (!longPressHandled) {

          if (pressDuration < LONG_PRESS_TIME) {

            // -----------------------------------
            // SHORT PRESS
            // -----------------------------------

            systemOn = !systemOn;


            if (systemOn) {

              Serial.println("SYSTEM ON");

              setLED(0, 255, 0);

            } else {

              Serial.println("SYSTEM OFF");

              calibrationMode = false;

              setLED(0, 0, 0);
            }
          }
        }
      }
    }
  }


  // =========================================
  // LONG PRESS DETECTION
  // =========================================

  if (buttonState == LOW &&
      !longPressHandled &&
      systemOn) {

    if ((millis() - buttonPressStart)
        >= LONG_PRESS_TIME) {

      longPressHandled = true;

      Serial.println("LONG PRESS DETECTED");

      resetCalibration();
    }
  }


  lastButtonReading = reading;
}


// =====================================================
// SETUP
// =====================================================

void setup() {

  Serial.begin(115200);

  Wire.begin();

  pinMode(BUTTON_PIN, INPUT_PULLUP);

  pinMode(RED_PIN, OUTPUT);
  pinMode(GREEN_PIN, OUTPUT);
  pinMode(BLUE_PIN, OUTPUT);

  setLED(0, 0, 0);

  setupMPU();

  delay(1000);

  Serial.println("Calibrating gyroscope...");

  calibrateGyro();

  Serial.println();
  Serial.println("=================================");
  Serial.println("SYSTEM READY");
  Serial.println("Short press = ON/OFF");
  Serial.println("Hold 2 seconds = CALIBRATE");
  Serial.println("=================================");

  previousTime =
    millis() / 1000.0;
}


// =====================================================
// MAIN LOOP
// =====================================================

void loop() {

  // Handle button first
  handleButton();


  // System is OFF
  if (!systemOn) {

    return;
  }


  // =============================================
  // READ MPU6050
  // =============================================

  int16_t ax, ay, az;
  int16_t gx, gy, gz;

  readMPU(
    ax, ay, az,
    gx, gy, gz
  );


  // =============================================
  // TIME
  // =============================================

  float currentTime =
    millis() / 1000.0;

  float dt =
    currentTime - previousTime;

  previousTime = currentTime;


  if (dt <= 0 || dt > 0.1) {

    dt = 0.01;
  }


  // =============================================
  // ACCELEROMETER ORIENTATION
  // =============================================

  float accelPitch =
    getAccelerometerPitch(
      ax, ay, az
    );


  float accelRoll =
    atan2(
      ay,
      sqrt(
        (float)ax * ax +
        (float)az * az
      )
    )
    * 180.0 / PI;

  accelRoll *= ROLL_SIGN;


  // =============================================
  // GYROSCOPE
  // =============================================

  float gyroX =
    (gx - gyroOffsetX) / 131.0;

  float gyroY =
    (gy - gyroOffsetY) / 131.0;

  float gyroZ =
    (gz - gyroOffsetZ) / 131.0;


  // =============================================
  // COMPLEMENTARY FILTER
  // =============================================

  pitch =
    0.98 * (pitch + gyroY * dt)
    + 0.02 * accelPitch;

  roll =
    0.98 * (roll + gyroX * dt)
    + 0.02 * accelRoll;


  // =============================================
  // MOVEMENT VARIABILITY
  // =============================================

  float deltaGX =
    gyroX - previousGyroX;

  float deltaGY =
    gyroY - previousGyroY;

  float deltaGZ =
    gyroZ - previousGyroZ;


  float shakeMagnitude =
    sqrt(
      deltaGX * deltaGX +
      deltaGY * deltaGY +
      deltaGZ * deltaGZ
    );


  liveShake =
    0.90 * liveShake
    + 0.10 * shakeMagnitude;


  previousGyroX = gyroX;
  previousGyroY = gyroY;
  previousGyroZ = gyroZ;


  // =============================================
  // REP DETECTION
  // =============================================

  if (calibrationMode) {

    detectCalibrationRep();

  } else if (calibrationComplete) {

    detectNormalRep();
  }


  // =============================================
  // LED FEEDBACK
  // =============================================

  updateFeedback();


  // =============================================
  // SERIAL MONITOR
  // =============================================

  static unsigned long lastPrint = 0;

  if (millis() - lastPrint > 250) {

    lastPrint = millis();

    Serial.print("Pitch: ");
    Serial.print(pitch);

    Serial.print(" | Bottom: ");
    Serial.print(baselineBottom);

    Serial.print(" | Top: ");
    Serial.print(baselineTop);

    Serial.print(" | Shake: ");
    Serial.print(liveShake);

    Serial.print(" | Calibration Reps: ");
    Serial.println(calibrationRepCount);
  }


  delay(5);
}