#include <WiFi.h>

// Replace these placeholders before uploading.
const char* ssid = "YOUR_SSID";
const char* password = "YOUR_PASSWORD";

// ESP32-S3 boards do not share a universal built-in LED pin. GPIO 2 is a
// conservative default for an external LED; change it for your board.
constexpr int LED_PIN = 2;
constexpr uint8_t LED_ON_LEVEL = HIGH;
constexpr uint8_t LED_OFF_LEVEL = LOW;

constexpr size_t WINDOW_SIZE = 40;
constexpr size_t LONG_WINDOW = 100;
constexpr uint16_t SAMPLE_INTERVAL_MS = 35;
constexpr uint16_t CALIBRATION_SAMPLES = 200;
constexpr uint32_t WIFI_CONNECT_TIMEOUT_MS = 30'000;

constexpr float PROCESS_NOISE = 0.15f;
constexpr float MEASUREMENT_NOISE = 0.3f;
constexpr float SMOOTH_ALPHA = 0.15f;
constexpr float ADAPTIVE_ALPHA = 0.008f;
constexpr float ADAPTIVE_BETA = 0.005f;
constexpr float SLOW_MOVEMENT_THRESHOLD = 1.8f;
constexpr float FAST_MOVEMENT_THRESHOLD = 3.5f;
constexpr float Z_SCORE_THRESHOLD = 2.5f;
constexpr float PEAK_THRESHOLD = 2.0f;
constexpr uint8_t PERSISTENCE_REQUIRED = 3;

float rssiWindow[WINDOW_SIZE];
float rssiLongWindow[LONG_WINDOW];
size_t windowIndex = 0;
size_t longWindowIndex = 0;
size_t sampleCount = 0;

float kalmanEstimate = -55.0f;
float kalmanError = 1.0f;
float smoothedRssi = -55.0f;
float baseline = -55.0f;
float baselineVariance = 1.0f;

uint8_t disturbanceCounter = 0;
uint32_t totalDetections = 0;

float clampValue(float value, float lower, float upper) {
  return value < lower ? lower : (value > upper ? upper : value);
}

size_t latestWindowIndex(size_t offset) {
  return (windowIndex + WINDOW_SIZE - offset) % WINDOW_SIZE;
}

float kalmanFilter(float measurement) {
  const float priorError = kalmanError + PROCESS_NOISE;
  const float gain = priorError / (priorError + MEASUREMENT_NOISE);

  kalmanEstimate += gain * (measurement - kalmanEstimate);
  kalmanError = (1.0f - gain) * priorError;
  return kalmanEstimate;
}

float exponentialSmoothing(float measurement) {
  smoothedRssi = SMOOTH_ALPHA * measurement + (1.0f - SMOOTH_ALPHA) * smoothedRssi;
  return smoothedRssi;
}

float standardDeviation(const float* values, size_t count) {
  if (count == 0) {
    return 0.0f;
  }

  float mean = 0.0f;
  for (size_t i = 0; i < count; ++i) {
    mean += values[i];
  }
  mean /= static_cast<float>(count);

  float variance = 0.0f;
  for (size_t i = 0; i < count; ++i) {
    const float difference = values[i] - mean;
    variance += difference * difference;
  }
  return sqrtf(variance / static_cast<float>(count));
}

float analyzeWindowVariance() {
  return standardDeviation(rssiWindow, WINDOW_SIZE);
}

float analyzeLongTermVariance() {
  return standardDeviation(rssiLongWindow, LONG_WINDOW);
}

float detectRateOfChange() {
  if (sampleCount < 10) {
    return 0.0f;
  }

  float recentMean = 0.0f;
  float oldMean = 0.0f;
  for (size_t i = 1; i <= 5; ++i) {
    recentMean += rssiWindow[latestWindowIndex(i)];
    oldMean += rssiWindow[latestWindowIndex(i + 5)];
  }
  return fabsf((recentMean - oldMean) / 5.0f);
}

float detectPeak() {
  if (sampleCount < 3) {
    return 0.0f;
  }

  const float current = rssiWindow[latestWindowIndex(1)];
  const float previous = rssiWindow[latestWindowIndex(2)];
  const float older = rssiWindow[latestWindowIndex(3)];
  const float currentDelta = fabsf(current - previous);
  const float previousDelta = fabsf(previous - older);

  return currentDelta > previousDelta + 1.0f ? currentDelta : 0.0f;
}

float calculateZScore() {
  if (sampleCount < WINDOW_SIZE) {
    return 0.0f;
  }

  float mean = 0.0f;
  for (size_t i = 0; i < WINDOW_SIZE; ++i) {
    mean += rssiWindow[i];
  }
  mean /= static_cast<float>(WINDOW_SIZE);

  const float stdDev = fmaxf(analyzeWindowVariance(), 0.1f);
  const float current = rssiWindow[latestWindowIndex(1)];
  return fabsf((current - mean) / stdDev);
}

int calculateConfidence(float variance, float rateOfChange, float peak, float zScore) {
  float confidence = variance > FAST_MOVEMENT_THRESHOLD ? 75.0f
                   : variance > SLOW_MOVEMENT_THRESHOLD ? 50.0f
                                                        : 15.0f;
  if (rateOfChange > 2.0f) confidence += 15.0f;
  if (peak > PEAK_THRESHOLD) confidence += 10.0f;
  if (zScore > Z_SCORE_THRESHOLD) confidence += 10.0f;
  return static_cast<int>(clampValue(confidence, 0.0f, 100.0f));
}

const char* getMotionIntensity(float variance, float rateOfChange) {
  if (variance > FAST_MOVEMENT_THRESHOLD && rateOfChange > 2.5f) return "SPRINT";
  if (variance > FAST_MOVEMENT_THRESHOLD) return "FAST";
  if (variance > SLOW_MOVEMENT_THRESHOLD && rateOfChange > 1.0f) return "WALKING";
  if (variance > SLOW_MOVEMENT_THRESHOLD) return "SLOW";
  return "CALM";
}

void initializeRssiWindows(float initialRssi) {
  for (int i = 0; i < WINDOW_SIZE; ++i) {
    rssiWindow[i] = initialRssi;
  }
  for (int i = 0; i < LONG_WINDOW; ++i) {
    rssiLongWindow[i] = initialRssi;
  }
  windowIndex = 0;
  longWindowIndex = 0;
  sampleCount = WINDOW_SIZE;
}

bool connectToWifi() {
  WiFi.mode(WIFI_STA);
  WiFi.begin(ssid, password);

  Serial.print("Connecting to Wi-Fi");
  const uint32_t startedAt = millis();
  while (WiFi.status() != WL_CONNECTED) {
    if (millis() - startedAt >= WIFI_CONNECT_TIMEOUT_MS) {
      Serial.println("\nWi-Fi connection timed out. Restarting...");
      return false;
    }
    delay(500);
    Serial.print('.');
  }

  Serial.printf("\nConnected. IP: %s, RSSI: %d dBm\n", WiFi.localIP().toString().c_str(), WiFi.RSSI());
  return true;
}

void calibrate() {
  Serial.printf("Calibrating from %u RSSI samples; keep the room empty...\n", CALIBRATION_SAMPLES);
  float calibrationMean = 0.0f;
  for (uint16_t i = 0; i < CALIBRATION_SAMPLES; ++i) {
    const float rssi = static_cast<float>(WiFi.RSSI());
    calibrationMean += rssi;
    delay(30);
  }

  baseline = calibrationMean / static_cast<float>(CALIBRATION_SAMPLES);
  kalmanEstimate = baseline;
  smoothedRssi = baseline;
  initializeRssiWindows(baseline);
  Serial.printf("Calibration complete. Baseline: %.2f dBm\n", baseline);
}

void setDetectionLed(bool detected) {
  digitalWrite(LED_PIN, detected ? LED_ON_LEVEL : LED_OFF_LEVEL);
}

void printMeasurements(
    int rawRssi,
    float kalmanFiltered,
    float smoothFiltered,
    float variance,
    float zScore,
    int confidence,
    float signalQuality,
    float rateOfChange) {
  Serial.printf(
      "Raw:%d,Kalman:%.2f,Smooth:%.2f,Baseline:%.2f,Variance:%.2f,ZScore:%.2f,Confidence:%d\n",
      rawRssi, kalmanFiltered, smoothFiltered, baseline, variance, zScore, confidence);
  Serial.printf(
      "[%s] Conf:%d%% | Quality:%d%% | Rate:%.1f | Total:%lu\n",
      getMotionIntensity(variance, rateOfChange), confidence, static_cast<int>(signalQuality), rateOfChange,
      static_cast<unsigned long>(totalDetections));
}

void setup() {
  Serial.begin(115200);
  pinMode(LED_PIN, OUTPUT);
  setDetectionLed(false);

  Serial.println("\n========== WiFiSense ESP32-S3 ==========");
  if (!connectToWifi()) {
    delay(2'000);
    ESP.restart();
  }
  calibrate();
  Serial.println("Open Tools > Serial Plotter to view RSSI metrics.");
}

void loop() {
  if (WiFi.status() != WL_CONNECTED) {
    setDetectionLed(false);
    if (!connectToWifi()) {
      delay(2'000);
    }
    return;
  }

  const int rawRssi = WiFi.RSSI();
  const float kalmanFiltered = kalmanFilter(static_cast<float>(rawRssi));
  const float smoothFiltered = exponentialSmoothing(kalmanFiltered);

  rssiWindow[windowIndex] = smoothFiltered;
  rssiLongWindow[longWindowIndex] = smoothFiltered;
  windowIndex = (windowIndex + 1) % WINDOW_SIZE;
  longWindowIndex = (longWindowIndex + 1) % LONG_WINDOW;
  ++sampleCount;

  const float variance = analyzeWindowVariance();
  const float longVariance = analyzeLongTermVariance();
  const float rateOfChange = detectRateOfChange();
  const float peak = detectPeak();
  const float zScore = calculateZScore();
  const int confidence = calculateConfidence(variance, rateOfChange, peak, zScore);
  const float signalQuality = clampValue(100.0f - longVariance * 20.0f, 0.0f, 100.0f);

  baseline = baseline * (1.0f - ADAPTIVE_ALPHA) + smoothFiltered * ADAPTIVE_ALPHA;
  baselineVariance = baselineVariance * (1.0f - ADAPTIVE_BETA) + variance * ADAPTIVE_BETA;

  const bool motion = variance > SLOW_MOVEMENT_THRESHOLD ||
                      (zScore > Z_SCORE_THRESHOLD && peak > PEAK_THRESHOLD);
  disturbanceCounter = motion ? disturbanceCounter + 1 : 0;

  const bool detected = disturbanceCounter >= PERSISTENCE_REQUIRED;
  if (detected && disturbanceCounter == PERSISTENCE_REQUIRED) {
    ++totalDetections;
  }
  setDetectionLed(detected);

  printMeasurements(rawRssi, kalmanFiltered, smoothFiltered, variance, zScore, confidence, signalQuality, rateOfChange);
  delay(SAMPLE_INTERVAL_MS);
}