#include <Arduino.h>
#include <Wire.h>
#include <LSM6DSL.h>           // dycodex/LSM6DSL-Arduino
#include <Adafruit_NeoPixel.h> // for the on-board WS2812
 
// ===== Sparrow pinout (from board README) =====
// I2C: SDA=GPIO21, SCL=GPIO22
// LSM6DSL I2C address = 0x6A
// NeoPixel on GPIO3
// (INT pin for LSM6DSL is GPIO11 if you later want interrupts)
constexpr int PIN_SDA = 21;
constexpr int PIN_SCL = 22;
constexpr uint8_t LSM6DSL_ADDR = 0x6A;
constexpr int NEOPIXEL_PIN = 3;
constexpr int NEOPIXEL_COUNT = 1;
 
Adafruit_NeoPixel pixel(NEOPIXEL_COUNT, NEOPIXEL_PIN, NEO_GRB + NEO_KHZ800);
LSM6DSL imu(LSM6DSL_ADDR);
 
// ===== Sampling & features =====
static const float G_SI = 9.80665f;
static const int SAMPLE_RATE_HZ = 104;        // ODR we'll ask the IMU to use
static const int WINDOW_MS = 2000;            // 2 s windows
static const int WINDOW_SAMPLES = SAMPLE_RATE_HZ * WINDOW_MS / 1000;
 
// We'll compute simple 4-D features per window:
//  [ mean(|a|), std(|a|), mean(|g|), std(|g|) ]
struct Feature {
  float f[4];
};
 
// ===== Simple K-means (from scratch) =====
template<int K>
struct KMeans {
  Feature centroids[K];
  bool initialized = false;
 
  static float dist2(const Feature& a, const Feature& b) {
    float d=0;
    for (int i=0;i<4;i++){ float t=a.f[i]-b.f[i]; d+=t*t; }
    return d;
  }
 
  int assign(const Feature& x) const {
    int best = 0; float bestd = dist2(x, centroids[0]);
    for (int k=1;k<K;k++){ float d=dist2(x, centroids[k]); if (d<bestd){best=k; bestd=d;} }
    return best;
  }
 
  // k-means++ init from a small set of bootstrap features
  void init_plus_plus(const Feature *bootstrap, int n) {
    // pick first centroid
    centroids[0] = bootstrap[random(n)];
    // next centroids
    for (int c=1;c<K;c++){
      // compute D(x)^2
      float sumD=0; static float D[128]; // bootstrap up to 128 windows
      int cap = min(n, 128);
      for (int i=0;i<cap;i++){
        float bestd = dist2(bootstrap[i], centroids[0]);
        for (int j=1;j<c;j++) bestd = min(bestd, dist2(bootstrap[i], centroids[j]));
        D[i]=bestd; sumD += bestd;
      }
      float r = random(0, 1000000) / 1000000.0f * sumD;
      int pick=0; float acc=0;
      for (; pick<cap; pick++){ acc += D[pick]; if (acc>=r) break; }
      centroids[c] = bootstrap[ pick < cap ? pick : cap-1 ];
    }
    initialized = true;
  }
 
  // batch k-means iterations
  void fit(const Feature *X, int n, int iters=10) {
    if (!initialized) { // fallback random init
      for (int k=0;k<K;k++) centroids[k] = X[random(n)];
      initialized = true;
    }
    int assignBuf[128];
    int cap = min(n, 128);
 
    for (int it=0; it<iters; it++){
      // assign
      for (int i=0;i<cap;i++) assignBuf[i]=assign(X[i]);
      // recompute
      Feature sums[K]; int counts[K]; 
      for (int k=0;k<K;k++){ for(int j=0;j<4;j++) sums[k].f[j]=0; counts[k]=0; }
      for (int i=0;i<cap;i++){
        int k = assignBuf[i];
        for (int j=0;j<4;j++) sums[k].f[j]+=X[i].f[j];
        counts[k]++;
      }
      for (int k=0;k<K;k++){
        if (counts[k]>0){
          for (int j=0;j<4;j++) centroids[k].f[j]=sums[k].f[j]/counts[k];
        }
      }
    }
  }
};
 
KMeans<3> kmeans;                   // K = 3 clusters
static Feature bootstrap[128];      // store first ~128 windows
static int bootstrap_count = 0;
static bool model_ready = false;
static bool stage_prompted[3] = {false, false, false};
constexpr int BOOTSTRAP_STAGE_WINDOWS = 10; // ~20 s total per instruction phase
 
// ===== helpers =====
void neopixel_show_cluster(int c){
  uint8_t r=0,g=0,b=0;
  if (c==0){ r=255; }          // red
  else if (c==1){ g=255; }     // green
  else { b=255; }              // blue
  pixel.setPixelColor(0, pixel.Color(r,g,b));
  pixel.show();
}
 
void print_feature(const Feature& x){
  Serial.printf("feat: aMean=%.3f, aStd=%.3f, gMean=%.3f, gStd=%.3f\n",
                x.f[0], x.f[1], x.f[2], x.f[3]);
}
 
void prompt_bootstrap_stage() {
  if (model_ready) {
    return;
  }
 
  if (bootstrap_count < BOOTSTRAP_STAGE_WINDOWS && !stage_prompted[0]) {
    Serial.println("Stage 1/3: move the board UP and DOWN for the first cluster.");
    Serial.println("Collecting motion windows... keep moving!");
    stage_prompted[0] = true;
  } else if (bootstrap_count >= BOOTSTRAP_STAGE_WINDOWS &&
             bootstrap_count < 2 * BOOTSTRAP_STAGE_WINDOWS && !stage_prompted[1]) {
    Serial.println("Stage 2/3: move the board LEFT and RIGHT for the second cluster.");
    Serial.println("Keep sweeping side-to-side until prompted again.");
    stage_prompted[1] = true;
  } else if (bootstrap_count >= 2 * BOOTSTRAP_STAGE_WINDOWS &&
             bootstrap_count < 3 * BOOTSTRAP_STAGE_WINDOWS && !stage_prompted[2]) {
    Serial.println("Stage 3/3: hold the board STILL to capture the stationary cluster.");
    Serial.println("Stay steady while we finish collecting samples.");
    stage_prompted[2] = true;
  }
}
 
// Compute features for one window (blocking read loop)
Feature compute_window_features() {
  // Welford variance on |a| and |g|
  double meanA=0, M2A=0, meanG=0, M2G=0;
  for (int i=1;i<=WINDOW_SAMPLES;i++){
    // read sensors
    // Library returns accel in g and gyro in deg/s based on configured ranges
    float aX = imu.readFloatAccelX() * G_SI;           // g -> m/s^2
    float aY = imu.readFloatAccelY() * G_SI;
    float aZ = imu.readFloatAccelZ() * G_SI;
    float gX = imu.readFloatGyroX() * DEG_TO_RAD;      // deg/s -> rad/s
    float gY = imu.readFloatGyroY() * DEG_TO_RAD;
    float gZ = imu.readFloatGyroZ() * DEG_TO_RAD;
 
    float amag = sqrtf(aX*aX + aY*aY + aZ*aZ);
    float gmag = sqrtf(gX*gX + gY*gY + gZ*gZ);
 
    // Welford updates
    double deltaA = amag - meanA; meanA += deltaA / i; M2A += deltaA*(amag - meanA);
    double deltaG = gmag - meanG; meanG += deltaG / i; M2G += deltaG*(gmag - meanG);
 
    delayMicroseconds(1000000 / SAMPLE_RATE_HZ); // crude pacing ~1/ODR
  }
  Feature x;
  x.f[0] = (float)meanA;
  x.f[1] = (float)sqrt(M2A / (WINDOW_SAMPLES-1));
  x.f[2] = (float)meanG;
  x.f[3] = (float)sqrt(M2G / (WINDOW_SAMPLES-1));
  return x;
}
 
void setupIMU() {
  // Configure desired ranges and ODR before starting the sensor
  imu.settings.accelRange = 2;       // ±2g
  imu.settings.accelSampleRate = 104;
  imu.settings.gyroRange = 245;      // ±245 dps
  imu.settings.gyroSampleRate = 104;
 
  if (imu.begin() != IMU_SUCCESS) {
    Serial.println("LSM6DSL not found. Check wiring/address.");
    while (1) { delay(1000); }
  }
}
 
void setup() {
  Serial.begin(115200);
  delay(200);
 
  // I2C
  Wire.begin(PIN_SDA, PIN_SCL, 400000); // fast-mode for quicker reads
 
  // NeoPixel
  pixel.begin();
  pixel.clear(); pixel.show();
 
  // IMU
  setupIMU();
 
  Serial.println("IMU Activity Clustering (K-means, K=3).");
  Serial.println("Bootstrapping ~80 s of guided motion for unsupervised init...");
  prompt_bootstrap_stage();
}
 
void loop() {
  prompt_bootstrap_stage();
 
  Feature feat = compute_window_features();
  print_feature(feat);
 
  if (!model_ready) {
    // collect bootstrap windows (cap 128)
    if (bootstrap_count < 128) {
      bootstrap[bootstrap_count++] = feat;
    }
    // After ~60 seconds (30 windows @ 2s) kick off training
    if (bootstrap_count >= 30 && !kmeans.initialized) {
      kmeans.init_plus_plus(bootstrap, bootstrap_count);
      kmeans.fit(bootstrap, bootstrap_count, /*iters=*/12);
      Serial.println("K-means initialized.");
    }
    // Consider model ready once we have a decent bootstrap
    if (bootstrap_count >= 40) {
      model_ready = true;
      Serial.println("Model ready. Streaming cluster assignments...");
    }
  } else {
    int c = kmeans.assign(feat);
    neopixel_show_cluster(c);
    Serial.printf("cluster: %d\n", c);
    // Optional: small online centroid update (streaming k-means)
    // move centroid slightly toward new sample
    const float alpha = 0.05f;
    for (int j=0;j<4;j++) {
      kmeans.centroids[c].f[j] = (1.0f - alpha) * kmeans.centroids[c].f[j] + alpha * feat.f[j];
    }
  }
}