// =============================================================================
// SpineTrack — Wokwi Simulation Test Sketch  (v2.1)
// Self-contained: No WiFi, No MQTT needed
// Tests: MPU6050 I2C read, complementary filter, 2-button debounce,
//        4-mode exercise rep detection FSM
// =============================================================================
// HOW TO USE IN WOKWI:
//   1. Go to https://wokwi.com/projects/new/esp32-s3
//   2. Replace diagram.json with contents of firmware/wokwi/diagram.json
//   3. Paste THIS entire file into the code editor tab
//   4. Click ▶ Play  →  open Serial Monitor @ 115200 baud
//   5. GREEN button (GPIO4) = Calibrate  |  BLUE button (GPIO5) = Switch Mode
// =============================================================================

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

// ── Pin Config ───────────────────────────────────────────────────────────────
#define PIN_SDA        8
#define PIN_SCL        9
#define PIN_BTN_CALIB  4    // GREEN button → GND (INPUT_PULLUP)
#define PIN_BTN_MODE   5    // BLUE  button → GND (INPUT_PULLUP)
#define PIN_LED        2

// ── MPU6050 Registers ────────────────────────────────────────────────────────
#define MPU6050_ADDR     0x68
#define MPU_PWR_MGMT_1   0x6B
#define MPU_ACCEL_XOUT_H 0x3B
#define ACCEL_SCALE      16384.0f   // ±2g
#define GYRO_SCALE       131.0f    // ±250°/s
#define LPF_ALPHA        0.15f     // smooth new data in (lower = more lag)

// ── Modes ────────────────────────────────────────────────────────────────────
#define MODE_POSTURE 0
#define MODE_PUSHUP  1
#define MODE_CRUNCH  2
#define MODE_SQUAT   3
#define MODE_COUNT   4
const char* MODE_NAMES[MODE_COUNT] = {"posture","pushup","crunch","squat"};

// ── Thresholds (degrees relative to baseline) ────────────────────────────────
// Pushup:  device on back prone, chin to floor = pitch decreases
const float PUSHUP_THRESHOLD_DOWN = -18.0f;
const float PUSHUP_THRESHOLD_UP   =  -4.0f;
// Crunch:  supine, torso curl = pitch increases
const float CRUNCH_THRESHOLD_UP   =  15.0f;
const float CRUNCH_THRESHOLD_DOWN =   6.0f;
// Squat:   standing, hip hinge forward = pitch increases
const float SQUAT_THRESHOLD_DOWN  =  22.0f;
const float SQUAT_THRESHOLD_UP    =   8.0f;
// Posture: deviation from baseline
const float POSTURE_SLOUCH_DEG    =   8.0f;

// ── State ─────────────────────────────────────────────────────────────────────
float pitch = 0.0f, roll = 0.0f;
float baseline_pitch = 0.0f;
bool  calibrated  = false;
bool  calibrating = false;
uint8_t mode  = MODE_POSTURE;
uint32_t reps = 0;
bool  in_rep  = false;
uint32_t slouch_count = 0;
bool was_slouching = false;

// Calibration accumulator
unsigned long calib_start_ms = 0;
float pitch_sum    = 0.0f;
int   pitch_samples = 0;

// Button state (non-blocking debounce)
bool last_calib = HIGH, last_mode = HIGH;
unsigned long calib_press_ms = 0, mode_press_ms = 0;
#define DEBOUNCE_MS 50

// FSM rep debouncer
unsigned long last_rep_ms = 0;
#define REP_DEBOUNCE_MS 800
#define NOISE_FLOOR_DEG 5.0f

// ── I2C helpers ───────────────────────────────────────────────────────────────
void mpu_write(uint8_t reg, uint8_t val) {
  Wire.beginTransmission(MPU6050_ADDR);
  Wire.write(reg); Wire.write(val);
  Wire.endTransmission();
}

bool mpu_read(int16_t &ax, int16_t &ay, int16_t &az,
              int16_t &gx, int16_t &gy, int16_t &gz) {
  Wire.beginTransmission(MPU6050_ADDR);
  Wire.write(MPU_ACCEL_XOUT_H);
  if (Wire.endTransmission(false) != 0) return false;
  Wire.requestFrom((uint8_t)MPU6050_ADDR, (uint8_t)14);
  if (Wire.available() < 14) return false;
  uint8_t b[14];
  for (int i = 0; i < 14; i++) b[i] = Wire.read();
  ax=(b[0]<<8)|b[1]; ay=(b[2]<<8)|b[3]; az=(b[4]<<8)|b[5];
  gx=(b[8]<<8)|b[9]; gy=(b[10]<<8)|b[11]; gz=(b[12]<<8)|b[13];
  return true;
}

// ── Setup ─────────────────────────────────────────────────────────────────────
void setup() {
  Serial.begin(115200);
  delay(200);

  Wire.begin(PIN_SDA, PIN_SCL);
  Wire.setClock(100000);  // 100 kHz — stable on Wokwi simulation
  mpu_write(MPU_PWR_MGMT_1, 0x00);  // wake MPU6050
  delay(150);

  pinMode(PIN_BTN_CALIB, INPUT_PULLUP);
  pinMode(PIN_BTN_MODE,  INPUT_PULLUP);
  pinMode(PIN_LED, OUTPUT);

  Serial.println(F("============================================"));
  Serial.println(F(" SpineTrack v2.1 — Wokwi Simulation"));
  Serial.println(F("============================================"));
  Serial.println(F(" GREEN btn (GPIO4) = Calibrate"));
  Serial.println(F(" BLUE  btn (GPIO5) = Switch Mode"));
  Serial.println(F(" Tilt MPU6050 widget to change pitch"));
  Serial.println(F("--------------------------------------------"));
  Serial.printf(" Mode: %s | Calibrated: NO\n", MODE_NAMES[mode]);
  Serial.println(F("============================================"));
}

// ── Main loop (50 Hz) ─────────────────────────────────────────────────────────
unsigned long last_loop_ms = 0;
unsigned long last_print_ms = 0;

void loop() {
  unsigned long now = millis();
  if (now - last_loop_ms < 20) return;   // 50 Hz
  last_loop_ms = now;

  // ── Read MPU6050 ──────────────────────────────────────────────────────────
  int16_t ax, ay, az, gx, gy, gz;
  if (mpu_read(ax, ay, az, gx, gy, gz)) {
    float fax = ax / ACCEL_SCALE;
    float fay = ay / ACCEL_SCALE;
    float faz = az / ACCEL_SCALE;
    // Accel-only pitch/roll (spec formula) + LPF
    float ap = atan2f(fax, sqrtf(fay*fay + faz*faz)) * (180.0f / M_PI);
    float ar = atan2f(fay, sqrtf(fax*fax + faz*faz)) * (180.0f / M_PI);
    pitch = (1.0f - LPF_ALPHA) * pitch + LPF_ALPHA * ap;
    roll  = (1.0f - LPF_ALPHA) * roll  + LPF_ALPHA * ar;
  }

  // ── Button: CALIB (GPIO4) ─────────────────────────────────────────────────
  bool cb = digitalRead(PIN_BTN_CALIB);
  if (cb == LOW && last_calib == HIGH) calib_press_ms = now;
  if (cb == HIGH && last_calib == LOW) {
    if (now - calib_press_ms > DEBOUNCE_MS && !calibrating) {
      calibrating    = true;
      calib_start_ms = now;
      pitch_sum      = 0.0f;
      pitch_samples  = 0;
      calibrated     = false;
      reps           = 0;
      in_rep         = false;
      Serial.printf("[CALIB] Started — hold STILL for 3s... (mode=%s)\n", MODE_NAMES[mode]);
    }
  }
  last_calib = cb;

  // ── Button: MODE (GPIO5) ──────────────────────────────────────────────────
  bool mb = digitalRead(PIN_BTN_MODE);
  if (mb == LOW && last_mode == HIGH) mode_press_ms = now;
  if (mb == HIGH && last_mode == LOW) {
    if (now - mode_press_ms > DEBOUNCE_MS) {
      mode      = (mode + 1) % MODE_COUNT;
      reps      = 0;
      in_rep    = false;
      calibrated = false;
      Serial.printf("[MODE] ► Switched to: %s\n", MODE_NAMES[mode]);
      Serial.println(F("[MODE]   Press CALIB button to re-calibrate for new mode"));
    }
  }
  last_mode = mb;

  // ── Calibration accumulation ──────────────────────────────────────────────
  if (calibrating) {
    pitch_sum += pitch;
    pitch_samples++;
    if (now - calib_start_ms >= 3000) {
      baseline_pitch = pitch_sum / pitch_samples;
      calibrated     = true;
      calibrating    = false;
      Serial.printf("[CALIB] ✓ Done! baseline=%.2f° (samples=%d)\n",
                    baseline_pitch, pitch_samples);
      Serial.println(F("[CALIB]   Sensor is now tracking — start exercising!"));
    }
  }

  // ── Exercise / Posture detection ──────────────────────────────────────────
  if (calibrated && !calibrating) {
    float rel = pitch - baseline_pitch;

    if (mode == MODE_POSTURE) {
      // Posture: slouch detection
      bool slouching = (fabsf(rel) > POSTURE_SLOUCH_DEG);
      if (slouching && !was_slouching) {
        slouch_count++;
        Serial.printf("[POSTURE] ⚠ Slouch detected! (rel=%.1f°) Count=%u\n",
                      rel, slouch_count);
      }
      was_slouching = slouching;

    } else if (fabsf(rel) > NOISE_FLOOR_DEG) {
      // Exercise rep FSM (3-state: IDLE → IN_REP → IDLE with debounce)
      unsigned long elapsed_rep = now - last_rep_ms;

      if (mode == MODE_PUSHUP) {
        if (!in_rep && rel <= PUSHUP_THRESHOLD_DOWN) {
          in_rep = true;
          Serial.printf("[PUSHUP] ↓ Down phase (rel=%.1f°)\n", rel);
        }
        if (in_rep && rel >= PUSHUP_THRESHOLD_UP && elapsed_rep >= REP_DEBOUNCE_MS) {
          in_rep = false; reps++; last_rep_ms = now;
          Serial.printf("[PUSHUP] ✓ Rep %u complete\n", reps);
        }

      } else if (mode == MODE_CRUNCH) {
        if (!in_rep && rel >= CRUNCH_THRESHOLD_UP) {
          in_rep = true;
          Serial.printf("[CRUNCH] ↑ Up phase (rel=%.1f°)\n", rel);
        }
        if (in_rep && rel <= CRUNCH_THRESHOLD_DOWN && elapsed_rep >= REP_DEBOUNCE_MS) {
          in_rep = false; reps++; last_rep_ms = now;
          Serial.printf("[CRUNCH] ✓ Rep %u complete\n", reps);
        }

      } else if (mode == MODE_SQUAT) {
        if (!in_rep && rel >= SQUAT_THRESHOLD_DOWN) {
          in_rep = true;
          Serial.printf("[SQUAT]  ↓ Hip hinge (rel=%.1f°)\n", rel);
        }
        if (in_rep && rel <= SQUAT_THRESHOLD_UP && elapsed_rep >= REP_DEBOUNCE_MS) {
          in_rep = false; reps++; last_rep_ms = now;
          Serial.printf("[SQUAT]  ✓ Rep %u complete\n", reps);
        }
      }
    }
  }

  // ── LED status ────────────────────────────────────────────────────────────
  if (calibrating) {
    digitalWrite(PIN_LED, (now / 100) % 2);   // fast blink = calibrating
  } else if (!calibrated) {
    digitalWrite(PIN_LED, (now / 1000) % 2);  // slow blink = idle
  } else {
    digitalWrite(PIN_LED, HIGH);              // solid = tracking
  }

  // ── Serial print @ 5 Hz ───────────────────────────────────────────────────
  if (now - last_print_ms >= 200) {
    last_print_ms = now;
    float rel = pitch - baseline_pitch;
    float score = 100.0f - min(100.0f, fabsf(rel) * 2.5f);
    Serial.printf(
      "pitch=%6.1f° | rel=%+6.1f° | mode=%-7s | calib=%d | reps=%u | inRep=%d | score=%.0f\n",
      pitch, rel, MODE_NAMES[mode], (int)calibrated, reps, (int)in_rep, score
    );
  }
}