// ============================================================
//  GRATITUDE LOOM — ESP32 Code v4
//  With Beat Clustering + Verbal Control + State Machine
//
//  THE KEY INSIGHT IN THIS VERSION:
//  Different weavers beat the weft differently.
//  Some beat once per weft pass. Some beat twice. Some three times.
//  That personal style is NOT inconsistency — it IS their rhythm.
//
//  What actually matters is the time between complete weaving
//  gestures (weft cycles), not between individual hammer strikes.
//
//  This version introduces BEAT CLUSTERING:
//  Beats that happen close together (within 0.6 seconds) are
//  treated as ONE weaving gesture — one weft cycle.
//  Rhythm and sustain are then measured cycle-to-cycle.
//
//  EXAMPLE:
//  Weaver beats twice per weft, every ~1 second:
//    beat-beat ......... beat-beat ......... beat-beat
//    [cluster 1]        [cluster 2]        [cluster 3]
//    Gap between clusters = ~1 second, BPM = 60
//    Gap within cluster   = ~0.2s, NOT counted as rhythm
//
//  VERBAL CONTROL (handled in P5.js, sent back via WebSocket):
//  "yes"              confirms transition when prompted
//  "stop/change/done" immediate transition any time
//  "no/stay"          cancels prompt, keeps weaving
//
//  LIBRARIES NEEDED (Tools > Manage Libraries):
//    MPU6050 by Electronic Cats
//    WiFiManager by tzapu
//    WebSockets by Markus Sattler
//    ArduinoJson by Benoit Blanchon
//    ESPmDNS built into ESP32 board package
//
//  WIRING:
//    MPU6050 VCC to ESP32 3.3V  (NOT 5V, damages sensor)
//    MPU6050 GND to ESP32 GND
//    MPU6050 SDA to ESP32 GPIO 21
//    MPU6050 SCL to ESP32 GPIO 22
// ============================================================


#include <Wire.h>
#include <MPU6050.h>
#include <WiFiManager.h>
#include <ESPmDNS.h>
#include <WebSocketsServer.h>
#include <ArduinoJson.h>


MPU6050 mpu;
WebSocketsServer webSocket(81);

#define BEAT_HISTORY     8
#define STATE_COMMITTED  0
#define STATE_READY      1
#define STATE_COMPLETING 2


// ── BEAT DETECTION ──────────────────────────────────────────
// Handles the raw sensor signal and basic debouncing

float beatThreshold = 1.6;    // G-force that triggers a beat impact
                               // Resting sensor = ~1.0G (gravity)
                               // Set halfway between resting and your swing peak
                               // Watch TotalAccel in Serial Plotter to calibrate

bool  inBeat        = false;  // TRUE while acceleration is above threshold
                               // THE LOCK: prevents one swing counting as many beats
                               // Without this: loop() runs 20x/sec, one swing = 10 "beats"

float peakForce     = 0;      // Highest G-force while inBeat is true
                               // Tracks the true peak, not just the trigger moment

unsigned long lastBeatTime = 0; // Millisecond of last beat, used for clustering


// ── BEAT CLUSTERING ─────────────────────────────────────────
// Groups rapid successive beats into one weft cycle
//
// CONCEPT:
// A weaver who beats twice per weft does: beat-beat...beat-beat...beat-beat
// The gap WITHIN a gesture (~0.2s) is not rhythm
// The gap BETWEEN gestures (~1.0s) IS rhythm
//
// RULE:
// New beat within clusterWindow of last beat? Same gesture, add to cluster
// New beat after clusterWindow? New gesture, new cycle begins
//
// A cluster completes when no beat arrives for clusterWindow seconds
// That completion is what we count as one weft cycle

float clusterWindow      = 0.6;  // Seconds — beats closer than this = same gesture
                                  // Typical double-beat gap: 0.1-0.3s (inside cluster)
                                  // Typical between-gesture gap: 0.5-2.0s (outside)
                                  // CALIBRATE: if double-beats are slow, raise to 0.8
                                  // If singles are merging wrongly, lower to 0.4

bool  inCluster          = false; // TRUE while a gesture cluster is active
int   beatsInCluster     = 0;     // Count of strikes in current cluster
int   beatsPerCycle      = 1;     // Completed cluster's beat count (weaver's style)
float peakForceInCluster = 0;     // Strongest hit across whole cluster
unsigned long clusterStartTime = 0; // When current cluster began


// ── RHYTHM (CYCLE-BASED) ────────────────────────────────────
// All rhythm now measured gesture-to-gesture, not beat-to-beat
// A double-beater scores the same as a single-beater
// if their gestures are equally consistent

float cycleDurations[BEAT_HISTORY]; // Last 8 cycle gaps in seconds
int   cycleIndex         = 0;
int   cycleCount         = 0;
float lastCycleDuration  = 0;
unsigned long lastCycleStartTime = 0; // Start time of previous cycle

float sustainScore       = 0;  // 0.0-1.0 consistency of cycle-to-cycle timing


// ── STATE MACHINE ───────────────────────────────────────────

int   weavingState       = STATE_COMMITTED;
int   currentCycleCount  = 0;   // Cycles done in this section
int   minCycles          = 5;   // Minimum cycles before transition possible
                                 // Counts GESTURES not individual beats
float readyThreshold     = 0.65; // Sustain score needed to enter READY
float completingPause    = 3.0;  // Seconds of pause in READY triggers COMPLETING

unsigned long readyPauseStart = 0;
bool          pauseStarted    = false;
int           sectionCount    = 0;


// ── PAUSE ───────────────────────────────────────────────────

float         pauseDuration   = 0;
bool          isPaused        = false;
bool          isReturning     = false;
unsigned long lastAnyBeatTime = 0;


// ════════════════════════════════════════════════════════════
//  SETUP
// ════════════════════════════════════════════════════════════

void setup() {

  Serial.begin(115200);
  Serial.println("\n=== GRATITUDE LOOM v4 ===");
  Serial.print("Cluster window: "); Serial.print(clusterWindow); Serial.println("s");
  Serial.print("Min cycles: "); Serial.println(minCycles);

  Wire.begin();

  mpu.initialize();
  if (!mpu.testConnection()) {
    Serial.println("MPU6050 not found! Check wiring.");
    while (1);
  }
  Serial.println("IMU ready");

  WiFiManager wm;
  if (!wm.autoConnect("GratitudeLoom-Setup")) {
    delay(3000);
    ESP.restart();
  }
  Serial.print("WiFi connected. IP: ");
  Serial.println(WiFi.localIP());

  MDNS.begin("gratitudeloom");

  webSocket.begin();
  Serial.println("WebSocket open on port 81");
  Serial.println("========================\n");

}


// ════════════════════════════════════════════════════════════
//  LOOP — 20 times per second
// ════════════════════════════════════════════════════════════

void loop() {

  webSocket.loop();


  // ── READ AND CONVERT SENSOR ──────────────────────────────

  int16_t ax, ay, az, gx, gy, gz;
  mpu.getMotion6(&ax, &ay, &az, &gx, &gy, &gz);

  // Divide by sensitivity constants to get real units
  // Acceleration: 16384 raw = 1G at default +/-2G range
  // Rotation: 131 raw = 1 degree/second at default +/-250 range
  float fax = ax / 16384.0, fay = ay / 16384.0, faz = az / 16384.0;
  float fgx = gx / 131.0,   fgy = gy / 131.0,   fgz = gz / 131.0;

  // Total 3D movement: Pythagorean theorem on all three axes
  // Direction-independent — works however the sensor is mounted
  float totalAccel = sqrt(fax*fax + fay*fay + faz*faz);


  // ── BEAT DETECTION (individual impacts) ─────────────────
  // Detects each individual hammer strike
  // The inBeat lock ensures one swing = one detection event

  bool newBeatDetected = false;

  if (totalAccel > beatThreshold && !inBeat) {
    inBeat          = true;
    peakForce       = totalAccel;
    newBeatDetected = true;
    lastAnyBeatTime = millis();
  }

  // Keep updating peak while still in beat (may not have peaked yet)
  if (inBeat && totalAccel > peakForce) {
    peakForce = totalAccel;
  }

  // Beat ends when acceleration drops to half threshold
  // Half-threshold creates hysteresis: triggers at 1.6G, resets at 0.8G
  // This gap prevents false re-triggers near the threshold
  if (totalAccel < beatThreshold * 0.5) {
    inBeat = false;
  }


  // ── BEAT CLUSTERING ─────────────────────────────────────
  // Groups beats into weft cycles based on timing
  // Only runs when a new beat was just detected

  if (newBeatDetected) {

    unsigned long now = millis();

    // How long since the last beat?
    float timeSinceLastBeat = (lastBeatTime > 0)
      ? (now - lastBeatTime) / 1000.0
      : 999.0; // Large number on first beat = forces new cluster

    lastBeatTime = now;


    // DECISION: New cluster or same cluster?
    bool startNewCluster = !inCluster || (timeSinceLastBeat > clusterWindow);

    if (startNewCluster) {

      // ── COMPLETE THE PREVIOUS CLUSTER ───────────────
      // Before starting new, close out the old one
      // This is where we measure a completed weft cycle

      if (inCluster && lastCycleStartTime > 0) {

        // Time between START of last cycle and START of this one
        // This is the true gesture-to-gesture rhythm
        float cycleDuration = (now - lastCycleStartTime) / 1000.0;

        // Store in circular buffer (8 slots, overwrites oldest)
        cycleDurations[cycleIndex] = cycleDuration;
        cycleIndex = (cycleIndex + 1) % BEAT_HISTORY;
        if (cycleCount < BEAT_HISTORY) cycleCount++;
        lastCycleDuration = cycleDuration;

        // Save this cluster's beat count as the weaver's style
        beatsPerCycle      = beatsInCluster;
        peakForceInCluster = peakForce;

        // Count completed cycle in this section
        currentCycleCount++;

        // SUSTAIN SCORE — consistency of cycle timing
        // How similar are the gaps between gestures?
        if (cycleCount >= 2) {

          float sum = 0;
          for (int i = 0; i < cycleCount; i++) sum += cycleDurations[i];
          float avg = sum / cycleCount;

          float variance = 0;
          for (int i = 0; i < cycleCount; i++) {
            variance += abs(cycleDurations[i] - avg);
          }
          float avgVariance = variance / cycleCount;

          // Low variance = high score (consistent = flow)
          // Inverted so 1.0 = consistent, 0.0 = chaotic
          sustainScore = 1.0 - constrain(avgVariance / avg, 0, 1);
        }

        // Print cycle summary
        float bpmNow = lastCycleDuration > 0 ? 60.0 / lastCycleDuration : 0;
        Serial.print("Cycle #"); Serial.print(currentCycleCount);
        Serial.print(" | x"); Serial.print(beatsPerCycle);
        Serial.print(" beats | BPM: "); Serial.print(bpmNow, 1);
        Serial.print(" | Sustain: "); Serial.print(sustainScore, 2);
        Serial.print(" | State: ");
        Serial.println(weavingState == 0 ? "COMMITTED" :
                       weavingState == 1 ? "READY" : "COMPLETING");
      }

      // ── START NEW CLUSTER ────────────────────────────
      inCluster          = true;
      clusterStartTime   = now;
      lastCycleStartTime = now;
      beatsInCluster     = 1;

    } else {

      // ── ADD TO EXISTING CLUSTER ──────────────────────
      // Another beat in the same weaving gesture
      beatsInCluster++;
      Serial.print("  strike "); Serial.print(beatsInCluster);
      Serial.println(" in cluster");
    }
  }


  // CLUSTER TIMEOUT
  // When no new beat arrives for clusterWindow seconds, the gesture is done
  if (inCluster && lastBeatTime > 0) {
    float elapsed = (millis() - lastBeatTime) / 1000.0;
    if (elapsed > clusterWindow) {
      inCluster = false; // Ready for next gesture
    }
  }


  // Calculate BPM from cycle timing
  float cycleBPM = (lastCycleDuration > 0) ? (60.0 / lastCycleDuration) : 0;


  // ── STATE MACHINE ────────────────────────────────────────

  if (weavingState == STATE_COMMITTED) {
    if (currentCycleCount >= minCycles && sustainScore >= readyThreshold) {
      weavingState = STATE_READY;
      pauseStarted = false;
      Serial.println("-> READY");
    }
  }

  else if (weavingState == STATE_READY) {
    float timeSinceCycle = lastCycleStartTime > 0
      ? (millis() - lastCycleStartTime) / 1000.0 : 0;

    if (timeSinceCycle > 1.0 && !pauseStarted && lastCycleStartTime > 0) {
      pauseStarted    = true;
      readyPauseStart = millis();
    }

    if (pauseStarted) {
      if ((millis() - readyPauseStart) / 1000.0 >= completingPause) {
        weavingState = STATE_COMPLETING;
        Serial.println("-> COMPLETING");
      }
    }

    if (newBeatDetected && pauseStarted) {
      pauseStarted = false;
    }
  }

  else if (weavingState == STATE_COMPLETING) {
    if (newBeatDetected) {
      weavingState = STATE_READY;
      pauseStarted = false;
      Serial.println("-> READY (resumed)");
    }
  }


  // ── WEBSOCKET MESSAGES FROM P5.JS ───────────────────────
  // P5.js sends these based on what the weaver says:
  //   "confirmed" = yes / stop / change / done / next
  //   "stay"      = no / stay / keep going

  webSocket.onEvent([](uint8_t num, WStype_t type, uint8_t* payload, size_t length) {
    if (type == WStype_TEXT) {
      String msg = String((char*)payload);

      if (msg == "confirmed") {
        // Complete section, reset everything for next pattern
        sectionCount++;
        currentCycleCount = 0;
        cycleCount        = 0;
        cycleIndex        = 0;
        sustainScore      = 0;
        lastCycleDuration = 0;
        inCluster         = false;
        beatsInCluster    = 0;
        weavingState      = STATE_COMMITTED;
        Serial.print("Section "); Serial.print(sectionCount);
        Serial.println(" complete -> next pattern");
      }

      if (msg == "stay" && weavingState == STATE_COMPLETING) {
        weavingState = STATE_READY;
        pauseStarted = false;
        Serial.println("Staying in pattern -> READY");
      }
    }
  });


  // ── LONG PAUSE DETECTION ────────────────────────────────
  // 30+ seconds without any beat = yarn change, warp adjustment
  // Rhythm memory fully preserved during pauses

  if (lastAnyBeatTime > 0) {
    pauseDuration = (millis() - lastAnyBeatTime) / 1000.0;
  }

  if (pauseDuration > 30.0 && !isPaused && !isReturning) {
    isPaused = true;
    Serial.println("Long pause — loom remembers");
  }

  if (isPaused && newBeatDetected) {
    isPaused    = false;
    isReturning = true;
  }

  if (isReturning && currentCycleCount >= 3) {
    isReturning = false;
  }


  // ── SEND DATA TO P5.JS ───────────────────────────────────

  StaticJsonDocument<400> doc;

  // Raw movement
  doc["totalAccel"] = totalAccel;
  doc["ax"] = fax; doc["ay"] = fay; doc["az"] = faz;
  doc["gx"] = fgx; doc["gy"] = fgy; doc["gz"] = fgz;

  // Beat and cluster
  doc["beatForce"]            = peakForceInCluster;
  doc["beatsPerCycle"]        = beatsPerCycle;
  doc["inCluster"]            = inCluster;
  doc["beatsInCurrentCluster"] = beatsInCluster;

  // Rhythm (cycle-based)
  doc["cycleDuration"] = lastCycleDuration;
  doc["cycleBPM"]      = cycleBPM;
  doc["sustainScore"]  = sustainScore;
  doc["cycleCount"]    = cycleCount;

  // Section and state
  doc["weavingState"]      = weavingState;
  doc["currentCycleCount"] = currentCycleCount;
  doc["minCycles"]         = minCycles;
  doc["sectionCount"]      = sectionCount;

  // Pause
  doc["isPaused"]      = isPaused;
  doc["isReturning"]   = isReturning;
  doc["pauseDuration"] = pauseDuration;
  doc["timestamp"]     = millis();

  String out;
  serializeJson(doc, out);
  webSocket.broadcastTXT(out);


  // ── SERIAL PLOTTER ───────────────────────────────────────
  // CycleBPM = your true weaving tempo
  // Sustain  = cycle consistency (rises as you find flow)
  // BtsPerCycle = your personal style (1, 2, or 3)
  Serial.print("TotalAccel:");   Serial.print(totalAccel, 2);
  Serial.print(",CycleBPM:");    Serial.print(cycleBPM, 1);
  Serial.print(",Sustain:");     Serial.print(sustainScore, 2);
  Serial.print(",BtsPerCycle:"); Serial.println(beatsPerCycle);

  delay(50);

} // end loop()