#include <Arduino.h>
#include <Adafruit_NeoPixel.h>
#include <math.h>
 
#include "LSM6DSL.h"
#include "model.h"  // generated by train_model.py
 
LSM6DSL imu;
 
// LED setup (NeoPixel on GPIO3, 1 pixel)
#define NEOPIXEL_PIN   3
#define NEOPIXEL_COUNT 1
Adafruit_NeoPixel pixel(NEOPIXEL_COUNT, NEOPIXEL_PIN, NEO_GRB + NEO_KHZ800);
 
// sampling
const uint32_t SAMPLE_PERIOD_MS = 10; // ~100 Hz
uint32_t lastSampleMs = 0;
 
// ring buffer for accel magnitude
static float magBuf[WINDOW_SIZE];
static uint16_t magIndex = 0;
static uint16_t sampleCountSinceLastEval = 0;
 
float logistic(float x) {
    // sigmoid
    return 1.0f / (1.0f + expf(-x));
}
 
void setLED(bool shaking) {
    if (shaking) {
        // green
        pixel.setPixelColor(0, pixel.Color(0, 255, 0));
    } else {
        // off
        pixel.setPixelColor(0, pixel.Color(0, 0, 0));
    }
    pixel.show();
}
 
void computeFeatures(float &mean_mag, float &std_mag, float &p2p_mag) {
    // Use the WINDOW_SIZE most recent samples in magBuf,
    // assuming magIndex is the "next write" index.
    float tmp[WINDOW_SIZE];
    for (int i = 0; i < WINDOW_SIZE; i++) {
        // oldest -> newest reconstruction
        int idx = (magIndex + i) % WINDOW_SIZE;
        tmp[i] = magBuf[idx];
    }
 
    // mean
    float sum = 0.0f;
    float minv = tmp[0];
    float maxv = tmp[0];
    for (int i = 0; i < WINDOW_SIZE; i++) {
        float v = tmp[i];
        sum += v;
        if (v < minv) minv = v;
        if (v > maxv) maxv = v;
    }
    mean_mag = sum / WINDOW_SIZE;
 
    // stddev
    float varSum = 0.0f;
    for (int i = 0; i < WINDOW_SIZE; i++) {
        float d = tmp[i] - mean_mag;
        varSum += d * d;
    }
    std_mag = sqrtf(varSum / WINDOW_SIZE);
 
    // peak-to-peak
    p2p_mag = maxv - minv;
}
 
void setup() {
    Serial.begin(115200);
    delay(2000);
 
    if (!imu.begin()) {
        Serial.println("# ERROR: IMU init failed");
        while (true) { delay(1000); }
    }
 
    pixel.begin();
    pixel.show(); // init off
 
    // init magBuf
    for (int i = 0; i < WINDOW_SIZE; i++) {
        magBuf[i] = 1.0f; // ~1g stationary baseline
    }
 
    Serial.println("# Sparrow shake detector running.");
}
 
void loop() {
    uint32_t now = millis();
    if (now - lastSampleMs >= SAMPLE_PERIOD_MS) {
        lastSampleMs = now;
 
        // 1. read imu
        float ax, ay, az;
        if (imu.readAccelG(ax, ay, az)) {
            // magnitude
            float mag = sqrtf(ax*ax + ay*ay + az*az);
 
            // ring buffer write
            magBuf[magIndex] = mag;
            magIndex = (magIndex + 1) % WINDOW_SIZE;
 
            sampleCountSinceLastEval++;
 
            // 2. every STEP_SIZE samples, evaluate model
            if (sampleCountSinceLastEval >= STEP_SIZE) {
                sampleCountSinceLastEval = 0;
 
                float mean_mag, std_mag, p2p_mag;
                computeFeatures(mean_mag, std_mag, p2p_mag);
 
                // 3. logistic regression
                float decision =
                    ShakeModel::W0 +
                    ShakeModel::W1 * mean_mag +
                    ShakeModel::W2 * std_mag +
                    ShakeModel::W3 * p2p_mag;
 
                float prob_shake = logistic(decision);
 
                bool shaking = (prob_shake > 0.5f);
 
                // 4. act
                setLED(shaking);
 
                // debug
                Serial.print("mean=");
                Serial.print(mean_mag, 4);
                Serial.print(" std=");
                Serial.print(std_mag, 4);
                Serial.print(" p2p=");
                Serial.print(p2p_mag, 4);
                Serial.print(" prob=");
                Serial.print(prob_shake, 3);
                Serial.print(" -> ");
                Serial.println(shaking ? "SHAKE" : "IDLE");
            }
        }
    }
}