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

// ── Pins
#define PIN_SDA        8
#define PIN_SCL        9
#define PIN_BTN_CALIB  4
#define PIN_BTN_MODE   5
#define PIN_LED        2

// ── MPU6050
#define MPU6050_ADDR     0x68
#define MPU_PWR_MGMT_1   0x6B
#define MPU_ACCEL_XOUT_H 0x3B
#define ACCEL_SCALE      16384.0f
#define LPF_ALPHA        0.15f

// ── 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 (deg from calibrated baseline)
const float PUSHUP_DOWN  = -18.0f, PUSHUP_UP   =  -4.0f;
const float CRUNCH_UP    =  15.0f, CRUNCH_DOWN  =   6.0f;
const float SQUAT_DOWN   =  22.0f, SQUAT_UP     =   8.0f;
const float SLOUCH_DEG   =   8.0f;
#define NOISE_FLOOR  5.0f
#define DEBOUNCE_MS  50
#define REP_DEB_MS   800

// ── State
float pitch = 0, roll = 0, baseline = 0;
bool calibrated = false, calibrating = false;
uint8_t mode = 0;
uint32_t reps = 0;
bool in_rep = false, was_slouching = false;
uint32_t slouch_count = 0;
unsigned long calib_start = 0;
float pitch_sum = 0;
int pitch_n = 0;
bool last_calib = HIGH, last_mode = HIGH;
unsigned long press_calib = 0, press_mode = 0, last_rep = 0;

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;
}

void setup() {
  Serial.begin(115200);
  delay(200);
  Wire.begin(PIN_SDA, PIN_SCL);
  Wire.setClock(100000);
  mpu_write(MPU_PWR_MGMT_1, 0x00);
  delay(150);
  pinMode(PIN_BTN_CALIB, INPUT_PULLUP);
  pinMode(PIN_BTN_MODE,  INPUT_PULLUP);
  pinMode(PIN_LED, OUTPUT);
  Serial.println("============================================");
  Serial.println(" SpineTrack v2.1 -- Wokwi Simulation");
  Serial.println("============================================");
  Serial.println(" GREEN btn (GPIO4) = Calibrate");
  Serial.println(" BLUE  btn (GPIO5) = Switch Mode");
  Serial.println(" Tilt MPU6050 widget to change pitch");
  Serial.println("--------------------------------------------");
  Serial.printf(" Mode: %s | Calibrated: NO\n", MODE_NAMES[mode]);
  Serial.println("============================================");
}

unsigned long t_loop = 0, t_print = 0;

void loop() {
  unsigned long now = millis();
  if (now - t_loop < 20) return;
  t_loop = now;

  // Read sensor
  int16_t ax,ay,az,gx,gy,gz;
  if (mpu_read(ax,ay,az,gx,gy,gz)) {
    float fx=ax/ACCEL_SCALE, fy=ay/ACCEL_SCALE, fz=az/ACCEL_SCALE;
    float ap = atan2f(fx,sqrtf(fy*fy+fz*fz))*(180.0f/M_PI);
    float ar = atan2f(fy,sqrtf(fx*fx+fz*fz))*(180.0f/M_PI);
    pitch = (1-LPF_ALPHA)*pitch + LPF_ALPHA*ap;
    roll  = (1-LPF_ALPHA)*roll  + LPF_ALPHA*ar;
  }

  // CALIB button
  bool cb = digitalRead(PIN_BTN_CALIB);
  if (cb==LOW && last_calib==HIGH) press_calib = now;
  if (cb==HIGH && last_calib==LOW && now-press_calib>DEBOUNCE_MS && !calibrating) {
    calibrating=true; calib_start=now; pitch_sum=0; pitch_n=0;
    calibrated=false; reps=0; in_rep=false;
    Serial.printf("[CALIB] Started -- hold STILL 3s... (mode=%s)\n", MODE_NAMES[mode]);
  }
  last_calib = cb;

  // MODE button
  bool mb = digitalRead(PIN_BTN_MODE);
  if (mb==LOW && last_mode==HIGH) press_mode = now;
  if (mb==HIGH && last_mode==LOW && now-press_mode>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("[MODE]   Press CALIB to re-calibrate for new mode");
  }
  last_mode = mb;

  // Calibration (3s average)
  if (calibrating) {
    pitch_sum += pitch; pitch_n++;
    if (now - calib_start >= 3000) {
      baseline=pitch_sum/pitch_n; calibrated=true; calibrating=false;
      Serial.printf("[CALIB] Done! baseline=%.2f deg (samples=%d)\n", baseline, pitch_n);
      Serial.println("[CALIB]   Tracking active -- start exercising!");
    }
  }

  // Detection
  if (calibrated && !calibrating) {
    float rel = pitch - baseline;
    if (mode == MODE_POSTURE) {
      bool sl = (fabsf(rel) > SLOUCH_DEG);
      if (sl && !was_slouching) { slouch_count++; Serial.printf("[POSTURE] Slouch! rel=%.1f Count=%u\n",rel,slouch_count); }
      was_slouching = sl;
    } else if (fabsf(rel) > NOISE_FLOOR) {
      unsigned long e = now - last_rep;
      if (mode==MODE_PUSHUP) {
        if (!in_rep && rel<=PUSHUP_DOWN)                      { in_rep=true;  Serial.printf("[PUSHUP] Down phase rel=%.1f\n",rel); }
        if ( in_rep && rel>=PUSHUP_UP   && e>=REP_DEB_MS)    { in_rep=false; reps++; last_rep=now; Serial.printf("[PUSHUP] Rep %u\n",reps); }
      } else if (mode==MODE_CRUNCH) {
        if (!in_rep && rel>=CRUNCH_UP)                        { in_rep=true;  Serial.printf("[CRUNCH] Up phase rel=%.1f\n",rel); }
        if ( in_rep && rel<=CRUNCH_DOWN && e>=REP_DEB_MS)    { in_rep=false; reps++; last_rep=now; Serial.printf("[CRUNCH] Rep %u\n",reps); }
      } else if (mode==MODE_SQUAT) {
        if (!in_rep && rel>=SQUAT_DOWN)                       { in_rep=true;  Serial.printf("[SQUAT] Hip hinge rel=%.1f\n",rel); }
        if ( in_rep && rel<=SQUAT_UP    && e>=REP_DEB_MS)    { in_rep=false; reps++; last_rep=now; Serial.printf("[SQUAT] Rep %u\n",reps); }
      }
    }
  }

  // LED
  if (calibrating)      digitalWrite(PIN_LED,(now/100)%2);
  else if (!calibrated) digitalWrite(PIN_LED,(now/1000)%2);
  else                  digitalWrite(PIN_LED,HIGH);

  // Status @ 5 Hz
  if (now - t_print >= 200) {
    t_print = now;
    float rel=pitch-baseline, 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);
  }
}