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

MPU6050 mpu(Wire);

float angle = 0.0;
float maxAngle = 0.0;
float downThreshold = 0.0;
float upThreshold = 0.0;
int repCount = 0;
bool goingDown = false;

int jointChoice = 0;
int exerciseChoice = 0;

const int resetButtonPin = 16;  // RESET button pin
unsigned long sessionStartTime = 0;

void setup() {
  Serial.begin(115200);
  Wire.begin();
  pinMode(resetButtonPin, INPUT_PULLUP);

  Serial.println("Initializing MPU6050...");
  byte status = mpu.begin();
  if (status == 0) Serial.println("MPU6050 ready!");
  else Serial.println("MPU6050 initialization failed!");

  // Start joint + exercise selection
  selectJointAndExercise();
}

void loop() {
  mpu.update();

  // ---- Static tilt angle using accelerometer ----
  // atan2(Y,Z) or atan2(X,Z) depending on axis; here using X/Z for demo
  float accX = mpu.getAccX();
  float accZ = mpu.getAccZ();
  angle = atan2(accX, accZ) * 180.0 / PI; // degrees
  if (angle < 0) angle = -angle;          // remove negatives

  // ---- Check within safe range ----
  if (angle < downThreshold || angle > upThreshold) {
    Serial.println("⚠ WARNING: Unsafe joint angle, risk of strain!");
  }

  // ---- Track max angle ----
  if (angle > maxAngle) maxAngle = angle;

  // ---- Rep counting logic ----
  if (!goingDown && angle < downThreshold + 2) { // small buffer
    goingDown = true;
  }
  if (goingDown && angle > upThreshold - 2) {
    repCount++;
    goingDown = false;
    Serial.print("Rep counted! Total reps: ");
    Serial.println(repCount);
  }

  // ---- Live angle display ----
  Serial.print("Current Angle: ");
  Serial.println(angle);

  // ---- Check if RESET button pressed (session end) ----
  if (digitalRead(resetButtonPin) == LOW) {
    delay(50); // debounce
    while (digitalRead(resetButtonPin) == LOW) {} // wait for release
    unsigned long totalDuration = millis() - sessionStartTime;
    printSessionSummary(totalDuration);
    selectJointAndExercise(); // return to joint selection
  }

  delay(200); // sampling rate
}

// ---------------- FUNCTIONS ---------------- //

void selectJointAndExercise() {
  // Reset session variables
  repCount = 0;
  maxAngle = 0;
  goingDown = false;

  sessionStartTime = millis();  // start timer for this session

  Serial.println("\n===== Physiotherapy Monitoring Menu =====");
  Serial.println("Select the joint you are monitoring:");
  Serial.println("1. Elbow");
  Serial.println("2. Knee");
  Serial.println("3. Ankle");
  Serial.println("4. Wrist");
  Serial.print("Enter option (1-4): ");
  while (Serial.available() == 0) {}
  jointChoice = Serial.parseInt();

  Serial.print("Selected joint: ");
  switch (jointChoice) {
    case 1: Serial.println("Elbow"); break;
    case 2: Serial.println("Knee"); break;
    case 3: Serial.println("Ankle"); break;
    case 4: Serial.println("Wrist"); break;
    default: Serial.println("Unknown joint"); break;
  }

  Serial.println("\nSelect the exercise for this joint:");
  switch (jointChoice) {
    case 1: Serial.println("1. Flexion\n2. Extension"); break;
    case 2: Serial.println("1. Flexion\n2. Extension"); break;
    case 3: Serial.println("1. Plantarflexion\n2. Dorsiflexion"); break;
    case 4: Serial.println("1. Flexion\n2. Extension\n3. Radial Deviation\n4. Ulnar Deviation"); break;
  }
  Serial.print("Enter option: ");
  while (Serial.available() == 0) {}
  exerciseChoice = Serial.parseInt();

  setThresholds(jointChoice, exerciseChoice);
  Serial.println("\nStart performing your exercise!");
}

void setThresholds(int joint, int exercise) {
  switch (joint) {
    case 1: // Elbow
      if (exercise == 1) { downThreshold = 5; upThreshold = 135; } // Flexion
      else if (exercise == 2) { downThreshold = 5; upThreshold = 135; } // Extension
      break;
    case 2: // Knee
      if (exercise == 1) { downThreshold = 5; upThreshold = 145; } // Flexion
      else if (exercise == 2) { downThreshold = 5; upThreshold = 145; } // Extension
      break;
    case 3: // Ankle
      if (exercise == 1) { downThreshold = 2; upThreshold = 38; } // Plantarflexion
      else if (exercise == 2) { downThreshold = 2; upThreshold = 28; } // Dorsiflexion
      break;
    case 4: // Wrist
      if (exercise == 1) { downThreshold = 5; upThreshold = 55; } // Flexion
      else if (exercise == 2) { downThreshold = 5; upThreshold = 55; } // Extension
      else if (exercise == 3) { downThreshold = 2; upThreshold = 18; } // Radial Dev.
      else if (exercise == 4) { downThreshold = 2; upThreshold = 18; } // Ulnar Dev.
      break;
  }
  Serial.print("Safe Down Threshold: "); Serial.println(downThreshold);
  Serial.print("Safe Up Threshold: "); Serial.println(upThreshold);
}

void printSessionSummary(unsigned long duration) {
  Serial.println("\n===== SESSION SUMMARY =====");
  Serial.println("{");
  Serial.print("  \"joint\": \"");
  switch (jointChoice) {
    case 1: Serial.print("Elbow"); break;
    case 2: Serial.print("Knee"); break;
    case 3: Serial.print("Ankle"); break;
    case 4: Serial.print("Wrist"); break;
  }
  Serial.println("\",");

  Serial.print("  \"exercise\": \"");
  switch (jointChoice) {
    case 1: Serial.print(exerciseChoice==1?"Flexion":"Extension"); break;
    case 2: Serial.print(exerciseChoice==1?"Flexion":"Extension"); break;
    case 3: Serial.print(exerciseChoice==1?"Plantarflexion":"Dorsiflexion"); break;
    case 4:
      if (exerciseChoice==1) Serial.print("Flexion");
      else if (exerciseChoice==2) Serial.print("Extension");
      else if (exerciseChoice==3) Serial.print("Radial Deviation");
      else if (exerciseChoice==4) Serial.print("Ulnar Deviation");
      break;
  }
  Serial.println("\",");

  Serial.print("  \"total_reps\": "); Serial.println(repCount);
  Serial.print("  \"max_angle\": "); Serial.println(maxAngle);

  float durationMinutes = duration / 60000.0; // convert ms to minutes
  Serial.print("  \"total_duration_minutes\": ");
  Serial.println(durationMinutes, 2);
  Serial.println("}");
}