// Dual Mode EMF Detector for ESP32 - SCANNER VERSION
// Mode 1: Classic EMF Detector
// Mode 2: WiFi Treasure Hunt (pure STA scan, no AP needed)

#include <WiFi.h>

// ========== PIN DEFINITIONS ==========
#define ANTENNA_PIN 15
#define MODE_SWITCH_PIN 4
#define MODE_LED_PIN 2

#define LED1 23
#define LED2 22
#define LED3 21
#define LED4 19
#define LED5 18

#define BUZZER_PIN 5

// ========== MODE 2 WIFI SETTINGS ==========
// This MUST match the ssid on your target ESP32(s)
const char* targetSsid = "emf_treasure_target";

// ========== MODE 1 EMF SETTINGS ==========
#define BUZZER_FREQ 800
#define THRESHOLD_1 100
#define THRESHOLD_2 200
#define THRESHOLD_3 400
#define THRESHOLD_4 700
#define THRESHOLD_5 1000

// ========== SCAN SETTINGS ==========
#define SCAN_INTERVAL 500  // ms between scans

// ========== GLOBAL VARIABLES ==========
// Mode 1
unsigned long lastBeepTime = 0;
int currentLevel = 0;
bool isBeeping = false;

// Mode 2
unsigned long lastWiFiBlinkTime = 0;
unsigned long lastScanTime = 0;
unsigned long lastBlinkTime = 0;
unsigned long beepOnTime = 0;
bool wifiLedState = false;
bool ledsOn = false;
int proximityLevel = 0;
int blinkTimer = 0;
bool targetsFound = false;
bool scanInProgress = false;

// Kalman filter state (one per target BSSID would be ideal,
// but since we only care about the closest, one global filter is fine)
float kalman_x = -65.0;  // Initial RSSI estimate (midrange guess)
float kalman_P = 1000.0; // Initial covariance — high = "I know nothing yet"
#define KALMAN_Q 2.0      // Process noise: how much the signal can legitimately change between scans
                          // Higher = filter reacts faster to real movement, but less smoothing
                          // Lower = smoother, but slower to react
#define KALMAN_R 8.0      // Measurement noise: how much you trust each raw RSSI reading
                          // Higher = trust raw readings less, smoother output
                          // Lower = trust raw readings more, noisier output

// Mode tracking
int currentMode = 1;
bool wifiInitialized = false;

void setup() {
  Serial.begin(115200);

  pinMode(MODE_SWITCH_PIN, INPUT_PULLUP);

  pinMode(LED1, OUTPUT);
  pinMode(LED2, OUTPUT);
  pinMode(LED3, OUTPUT);
  pinMode(LED4, OUTPUT);
  pinMode(LED5, OUTPUT);
  pinMode(MODE_LED_PIN, OUTPUT);

  ledcAttach(BUZZER_PIN, BUZZER_FREQ, 8);

  digitalWrite(LED1, LOW);
  digitalWrite(LED2, LOW);
  digitalWrite(LED3, LOW);
  digitalWrite(LED4, LOW);
  digitalWrite(LED5, LOW);
  digitalWrite(MODE_LED_PIN, LOW);

  Serial.println("Dual Mode EMF Detector Ready!");
  startupSequence();
}

void loop() {
  int switchState = digitalRead(MODE_SWITCH_PIN);
  currentMode = (switchState == LOW) ? 2 : 1;

  if (currentMode == 1) {
    runEMFMode();
  } else {
    runWiFiMode();
  }
}

// ========== MODE 1: EMF DETECTOR ==========
void runEMFMode() {
  if (wifiInitialized) {
    WiFi.scanDelete();
    WiFi.mode(WIFI_OFF);
    wifiInitialized = false;
    proximityLevel = 0;
    targetsFound = false;
    scanInProgress = false;
    digitalWrite(MODE_LED_PIN, LOW);
    digitalWrite(LED1, LOW);
    digitalWrite(LED2, LOW);
    digitalWrite(LED3, LOW);
    digitalWrite(LED4, LOW);
    digitalWrite(LED5, LOW);
    buzzerOff();
    Serial.println("Switched to EMF Mode");
  }

  int emfValue = analogRead(ANTENNA_PIN);

  static int smoothedValue = 0;
  smoothedValue = (smoothedValue * 9 + emfValue) / 10;

  int level = getLevel(smoothedValue);
  currentLevel = level;

  updateLEDs(level);
  handleBuzzer(level);

  Serial.print("EMF Mode | Raw: ");
  Serial.print(emfValue);
  Serial.print(" | Smoothed: ");
  Serial.print(smoothedValue);
  Serial.print(" | Level: ");
  Serial.println(level);

  delay(50);
}

int getLevel(int value) {
  if (value >= THRESHOLD_5) return 5;
  if (value >= THRESHOLD_4) return 4;
  if (value >= THRESHOLD_3) return 3;
  if (value >= THRESHOLD_2) return 2;
  if (value >= THRESHOLD_1) return 1;
  return 0;
}

void updateLEDs(int level) {
  digitalWrite(LED1, level >= 1 ? HIGH : LOW);
  digitalWrite(LED2, level >= 2 ? HIGH : LOW);
  digitalWrite(LED3, level >= 3 ? HIGH : LOW);
  digitalWrite(LED4, level >= 4 ? HIGH : LOW);
  digitalWrite(LED5, level >= 5 ? HIGH : LOW);
}

void handleBuzzer(int level) {
  unsigned long currentTime = millis();

  if (level == 0) {
    buzzerOff();
    lastBeepTime = currentTime;
    isBeeping = false;
    return;
  }

  if (level == 5) {
    buzzerOn();
    return;
  }

  int beepInterval = 1000 / level;

  if (!isBeeping && (currentTime - lastBeepTime >= beepInterval)) {
    buzzerOn();
    isBeeping = true;
    lastBeepTime = currentTime;
  } else if (isBeeping && (currentTime - lastBeepTime >= 50)) {
    buzzerOff();
    isBeeping = false;
    lastBeepTime = currentTime;
  }
}

// ========== MODE 2: WIFI TREASURE HUNT ==========
void runWiFiMode() {
  if (!wifiInitialized) {
    digitalWrite(LED1, LOW);
    digitalWrite(LED2, LOW);
    digitalWrite(LED3, LOW);
    digitalWrite(LED4, LOW);
    digitalWrite(LED5, LOW);
    buzzerOff();

    Serial.println("Switching to WiFi Hunt Mode (STA scan)");

    WiFi.mode(WIFI_STA);
    WiFi.disconnect();
    delay(100);

    wifiInitialized = true;
    lastScanTime = 0;
    scanInProgress = false;
    proximityLevel = 0;
    targetsFound = false;
    beepOnTime = 0;
    kalman_x = -65.0;
    kalman_P = 1000.0;
  }

  unsigned long now = millis();

  // Async scan state machine runs quietly in background
  if (!scanInProgress && (now - lastScanTime >= SCAN_INTERVAL)) {
    WiFi.scanNetworks(true);  // async, non-blocking
    scanInProgress = true;
    lastScanTime = now;
  }

  if (scanInProgress) {
    int n = WiFi.scanComplete();
    if (n != WIFI_SCAN_RUNNING) {
      scanInProgress = false;
      processScanResults(n);  // Only updates level if it actually changed
    }
  }

  // These run every single loop iteration, completely free from scan timing
  handleHuntLEDs();
  handleHuntBuzzer();
  handleWiFiStatusLED();
}

// Simple 1D Kalman filter — smooths noisy RSSI into a stable estimate
float kalmanFilter(float measurement) {
  // Predict: uncertainty grows a bit each step
  kalman_P = kalman_P + KALMAN_Q;

  // Update: how much do we trust this new measurement vs our current estimate?
  float K = kalman_P / (kalman_P + KALMAN_R);  // Kalman gain
  kalman_x = kalman_x + K * (measurement - kalman_x);  // Blend old estimate with new reading
  kalman_P = (1.0 - K) * kalman_P;  // Shrink uncertainty

  return kalman_x;
}

void processScanResults(int n) {
  if (n <= 0) {
    static int failCount = 0;
    failCount++;
    if (failCount >= 5) {
      Serial.print("Scan returned: ");
      Serial.println(n);
      failCount = 0;
    }

    if (targetsFound) {
      targetsFound = false;
      proximityLevel = 0;
      blinkTimer = 0;
      ledsOn = false;
      beepOnTime = 0;
      buzzerOff();
      // Reset Kalman state so it doesn't remember stale values
      kalman_x = -70.0;
      kalman_P = 1000.0;
      Serial.println("Target lost");
    }

    WiFi.scanDelete();
    return;
  }

  int bestRssi = -999;

  for (int i = 0; i < n; i++) {
    if (WiFi.SSID(i) == targetSsid) {
      int rssi = WiFi.RSSI(i);
      if (rssi > bestRssi) bestRssi = rssi;
    }
  }

  WiFi.scanDelete();

  if (bestRssi == -999) {
    if (targetsFound) {
      targetsFound = false;
      proximityLevel = 0;
      blinkTimer = 0;
      ledsOn = false;
      beepOnTime = 0;
      buzzerOff();
      kalman_x = -70.0;
      kalman_P = 1000.0;
      Serial.println("Target lost");
    }
    return;
  }

  targetsFound = true;

  // Run raw RSSI through Kalman filter
  float smoothedRssi = kalmanFilter((float)bestRssi);

  // Map the SMOOTHED value to levels, not the raw one
  int newLevel = 0;
  int newTimer = 0;

  if (smoothedRssi <= -80)      { newLevel = 0; newTimer = 0;    }
  else if (smoothedRssi <= -75) { newLevel = 1; newTimer = 2000; }
  else if (smoothedRssi <= -70) { newLevel = 2; newTimer = 1500; }
  else if (smoothedRssi <= -60) { newLevel = 3; newTimer = 1000; }
  else if (smoothedRssi <= -50) { newLevel = 4; newTimer = 500;  }
  else                          { newLevel = 5; newTimer = 0;    }

  if (newLevel != proximityLevel) {
    proximityLevel = newLevel;
    blinkTimer = newTimer;
    lastBlinkTime = millis();
    ledsOn = true;
    beepOnTime = 0;
    Serial.print("Raw: ");
    Serial.print(bestRssi);
    Serial.print(" | Kalman: ");
    Serial.print(smoothedRssi);
    Serial.print(" -> Level ");
    Serial.println(proximityLevel);
  }
}

void handleHuntLEDs() {
  unsigned long now = millis();

  if (proximityLevel == 0) {
    digitalWrite(LED1, LOW);
    digitalWrite(LED2, LOW);
    digitalWrite(LED3, LOW);
    digitalWrite(LED4, LOW);
    digitalWrite(LED5, LOW);
  } else if (proximityLevel == 5) {
    digitalWrite(LED1, HIGH);
    digitalWrite(LED2, HIGH);
    digitalWrite(LED3, HIGH);
    digitalWrite(LED4, HIGH);
    digitalWrite(LED5, HIGH);
  } else {
    // Runs freely every loop, totally independent of scan
    int halfCycle = blinkTimer / 2;

    if (now - lastBlinkTime >= halfCycle) {
      ledsOn = !ledsOn;
      lastBlinkTime = now;
    }

    digitalWrite(LED1, (proximityLevel >= 1 && ledsOn) ? HIGH : LOW);
    digitalWrite(LED2, (proximityLevel >= 2 && ledsOn) ? HIGH : LOW);
    digitalWrite(LED3, (proximityLevel >= 3 && ledsOn) ? HIGH : LOW);
    digitalWrite(LED4, (proximityLevel >= 4 && ledsOn) ? HIGH : LOW);
    digitalWrite(LED5, LOW);
  }
}

void handleHuntBuzzer() {
  unsigned long now = millis();

  if (proximityLevel == 0) {
    buzzerOff();
    beepOnTime = 0;
  } else if (proximityLevel == 5) {
    buzzerOn();
    beepOnTime = 0;
  } else {
    // Short 50ms beep fires at the start of each "LEDs on" phase
    if (ledsOn && beepOnTime == 0) {
      buzzerOn();
      beepOnTime = now;  // Start timing the beep
    }
    if (beepOnTime > 0 && (now - beepOnTime >= 50)) {
      buzzerOff();       // 50ms done, cut it
      beepOnTime = -1;   // Set to -1 so we don't re-trigger until next blink cycle
    }
    if (!ledsOn) {
      beepOnTime = 0;    // LEDs off = reset, ready to beep on next "on" phase
    }
  }
}

void handleWiFiStatusLED() {
  unsigned long now = millis();

  if (!targetsFound) {
    if (now - lastWiFiBlinkTime >= 125) {
      wifiLedState = !wifiLedState;
      digitalWrite(MODE_LED_PIN, wifiLedState);
      lastWiFiBlinkTime = now;
    }
  } else {
    digitalWrite(MODE_LED_PIN, HIGH);
  }
}

// ========== BUZZER CONTROL ==========
void buzzerOn() {
  ledcWrite(BUZZER_PIN, 128);
}

void buzzerOff() {
  ledcWrite(BUZZER_PIN, 0);
}

// ========== STARTUP SEQUENCE ==========
void startupSequence() {
  for (int i = 0; i < 2; i++) {
    digitalWrite(LED1, HIGH); delay(100);
    digitalWrite(LED2, HIGH); delay(100);
    digitalWrite(LED3, HIGH); delay(100);
    digitalWrite(LED4, HIGH); delay(100);
    digitalWrite(LED5, HIGH); delay(100);

    digitalWrite(LED1, LOW);
    digitalWrite(LED2, LOW);
    digitalWrite(LED3, LOW);
    digitalWrite(LED4, LOW);
    digitalWrite(LED5, LOW);
    delay(100);
  }

  buzzerOn();
  delay(50);
  buzzerOff();
  delay(50);
  buzzerOn();
  delay(50);
  buzzerOff();
  delay(50);
  buzzerOn();
  delay(250);
  buzzerOff();
}