# ESP32 Test

#include <Arduino.h>
#include <Wire.h>
#include <vl53l4cx_class.h>
#include <Adafruit_LSM6DSOX.h>
#include <Adafruit_LIS3MDL.h>
#include <Adafruit_AHRS.h>

//#define TOF_DEBUG 1
#define USE_XSHUT_PIN 1

// ==== Pin Assignments ====
#if USE_XSHUT_PIN
constexpr uint8_t kXshutPin   = 4;   // ToF XSHUT pin
#endif
constexpr uint8_t kSdaPin     = 19;  // Sensor I2C SDA
constexpr uint8_t kSclPin     = 18;  // Sensor I2C SCL
constexpr uint8_t kPiSdaPin   = 21;  // RPi-ESP32 I2C SDA
constexpr uint8_t kPiSclPin   = 22;  // RPi-ESP32 I2C SCL
constexpr uint8_t kSlaveAddr  = 0x54;// ESP32 I2C slave address for RPi

constexpr uint8_t kMotorPwmPin   = 25;
constexpr uint8_t kMotorDirPin   = 26;
constexpr uint8_t kMotorBrakePin = 27;
constexpr uint8_t kEncoderA_Pin  = 35;
constexpr uint8_t kEncoderB_Pin  = 34;

// ==== Constants ====
constexpr uint32_t kI2cFreqHz      = 400000;
constexpr uint32_t kSerialBaud     = 115200;
constexpr uint32_t kMotorPwmFreq   = 5000;
constexpr uint8_t  kMotorPwmRes    = 8;
constexpr float    kTicksPerRev    = 100.0;

// ==== ToF-LIDAR Globals ====
TwoWire *activeBus = &Wire;
VL53L4CX tof;

enum RangeMode : uint8_t {
  MODE_SHORT  = VL53L4CX_DISTANCEMODE_SHORT,
  MODE_MEDIUM = VL53L4CX_DISTANCEMODE_MEDIUM,
  MODE_LONG   = VL53L4CX_DISTANCEMODE_LONG
};
static RangeMode gCurrentMode = MODE_SHORT;
static uint8_t gNoTargetFrames     = 0;
static uint8_t gStableTargetFrames = 0;
constexpr uint8_t kNoTargetThreshold     = 15;
constexpr uint8_t kStableTargetThreshold = 8;
constexpr int16_t kShortMaxMm  = 1200;
constexpr int16_t kMediumMaxMm = 3500;
static int16_t gPrevDistanceMm = -1;

static uint8_t gConsecutiveReadErrors = 0;
constexpr uint8_t kReadErrorThreshold = 4;

// ==== I2C Slave Globals ====
volatile uint16_t gLatestDistanceMm = 0;
volatile bool     gHasValidDistance = false;
volatile unsigned long gLastPiRequestMs = 0;
constexpr uint16_t kNoTargetValue = 0xFFFF;

// ==== Encoder Globals ====
volatile long gEncoderTicks = 0;
unsigned long gLastRpsTime = 0;
long gLastRpsTicks = 0;

// ==== IMU + Sensor Fusion Globals ====
Adafruit_LSM6DSOX sox;
Adafruit_LIS3MDL lis3mdl;
Adafruit_NXPSensorFusion filter;
float gYaw = 0, gPitch = 0, gRoll = 0;

// ==== Function Prototypes ====
void onPiRequest();
void onPiReceive(int);
void maybeReportMissingPi();
void IRAM_ATTR updateEncoder();
float getRPS();
void brakeAndWait(float stopThresholdRPS = 0.2);

static void applyMode(RangeMode mode) {
  tof.VL53L4CX_Off();
  uint32_t budget_us = 8000;
  switch (mode) {
    case MODE_SHORT:  budget_us = 8000;  break;
    case MODE_MEDIUM: budget_us = 16000; break;
    case MODE_LONG:   budget_us = 33000; break;
  }
  tof.VL53L4CX_SetDistanceMode(static_cast<uint16_t>(mode));
  tof.VL53L4CX_SetMeasurementTimingBudgetMicroSeconds(budget_us);
  tof.VL53L4CX_StartMeasurement();
}

#if USE_XSHUT_PIN
static bool sensorHardReset() {
  digitalWrite(kXshutPin, LOW);
  delay(10);
  digitalWrite(kXshutPin, HIGH);
  delay(50);
  if (tof.InitSensor(VL53L4CX_DEFAULT_DEVICE_ADDRESS) != VL53L4CX_ERROR_NONE) {
    Serial.println(F("[ToF] Re-init after reset failed"));
    return false;
  }
  applyMode(gCurrentMode);
  Serial.println(F("[ToF] Sensor recovered via XSHUT reset"));
  return true;
}
#endif

static const __FlashStringHelper *statusText(uint8_t s) {
  switch (s) {
    case VL53L4CX_RANGESTATUS_RANGE_VALID:                     return F("VALID");
    case VL53L4CX_RANGESTATUS_SIGMA_FAIL:                      return F("SIGMA");
    case VL53L4CX_RANGESTATUS_OUTOFBOUNDS_FAIL:                return F("OOB");
    case VL53L4CX_RANGESTATUS_HARDWARE_FAIL:                   return F("HW_FAIL");
    case VL53L4CX_RANGESTATUS_RANGE_VALID_MIN_RANGE_CLIPPED:   return F("CLIPPED");
    default:                                                   return F("UNKNOWN");
  }
}

static bool probeDevice(uint8_t address, TwoWire &bus) {
  bus.beginTransmission(address);
  return (bus.endTransmission() == 0);
}

void onPiRequest() {
  uint16_t d = gHasValidDistance ? gLatestDistanceMm : kNoTargetValue;
  Wire1.write(static_cast<uint8_t>(d >> 8));
  Wire1.write(static_cast<uint8_t>(d & 0xFF));
  gLastPiRequestMs = millis();
}

void onPiReceive(int) {
  while (Wire1.available()) Wire1.read();
}

void maybeReportMissingPi() {
  static unsigned long lastWarn = 0;
  const unsigned long now = millis();
  const unsigned long kTimeoutMs = 5000;
  if (now - gLastPiRequestMs > kTimeoutMs && now - lastWarn > kTimeoutMs) {
    Serial.println(F("[I2C] No Pi requests detected – continuing measurements"));
    lastWarn = now;
  }
}

// ====== Motor & Encoder Functions =======
void IRAM_ATTR updateEncoder() {
  if (digitalRead(kEncoderB_Pin) == LOW) {
    gEncoderTicks++;
  } else {
    gEncoderTicks--;
  }
}

float getRPS() {
  unsigned long currentTime = millis();
  long currentTicks = gEncoderTicks;
  unsigned long deltaTime = currentTime - gLastRpsTime;
  if (deltaTime == 0) return 0.0;
  long deltaTicks = currentTicks - gLastRpsTicks;
  gLastRpsTime = currentTime;
  gLastRpsTicks = currentTicks;
  return (float)abs(deltaTicks) / kTicksPerRev / ((float)deltaTime / 1000.0);
}

void brakeAndWait(float stopThresholdRPS) {
  digitalWrite(kMotorDirPin, LOW);
  digitalWrite(kMotorBrakePin, LOW);
  ledcWrite(kMotorPwmPin, 0);
  delay(100);
  while (true) {
    if (getRPS() <= stopThresholdRPS) break;
    delay(50);
  }
  delay(100);
}

void setup() {
  Serial.begin(kSerialBaud);
  while (!Serial) { ; }
  Serial.println(F("Booting Combined ToF-LIDAR + IMU Controller..."));

  // --- Motor & Encoder Setup ---
  pinMode(kMotorDirPin, OUTPUT);
  pinMode(kMotorBrakePin, OUTPUT);
  pinMode(kEncoderA_Pin, INPUT_PULLUP);
  pinMode(kEncoderB_Pin, INPUT_PULLUP);
  digitalWrite(kMotorDirPin, HIGH); 
  digitalWrite(kMotorBrakePin, LOW);
  attachInterrupt(digitalPinToInterrupt(kEncoderA_Pin), updateEncoder, RISING);
  ledcAttach(kMotorPwmPin, kMotorPwmFreq, kMotorPwmRes);
  ledcWrite(kMotorPwmPin, 0);
  brakeAndWait();

  // --- ToF Sensor Setup ---
#if USE_XSHUT_PIN
  pinMode(kXshutPin, OUTPUT);
  digitalWrite(kXshutPin, LOW);
  delay(10);
  digitalWrite(kXshutPin, HIGH);
  delay(50);
#endif

  Wire.begin(kSdaPin, kSclPin);
  Wire.setClock(kI2cFreqHz);

  // --- ToF init ---
  if (!probeDevice(VL53L4CX_DEFAULT_DEVICE_ADDRESS >> 1, Wire)) {
    Serial.println(F("[ToF] Sensor not detected on primary I2C bus. Halting."));
    while (true) delay(1000);
  }
  tof.setI2cDevice(activeBus);
  if (tof.InitSensor(VL53L4CX_DEFAULT_DEVICE_ADDRESS) != VL53L4CX_ERROR_NONE) {
    Serial.println(F("[ToF] InitSensor failed. Halting."));
    while (true) delay(1000);
  }
  applyMode(gCurrentMode);

  // --- IMU Init (LSM6DSOX + LIS3MDL) ---
  if (!sox.begin_I2C(0x6A, &Wire)) {
    Serial.println("Failed to find LSM6DSOX chip");
    while (1) delay(10);
  }
  if (!lis3mdl.begin_I2C(0x1C, &Wire)) {
    Serial.println("Failed to find LIS3MDL chip");
    while (1) delay(10);
  }
  sox.setAccelRange(LSM6DS_ACCEL_RANGE_4_G);
  sox.setGyroRange(LSM6DS_GYRO_RANGE_500_DPS);
  lis3mdl.setRange(LIS3MDL_RANGE_4_GAUSS);
  filter.begin(200); // Sensor fusion update rate (Hz)

  // --- RPi I2C Slave Setup ---
  Wire1.begin(kSlaveAddr, kPiSdaPin, kPiSclPin, kI2cFreqHz);
  Wire1.onRequest(onPiRequest);
  Wire1.onReceive(onPiReceive);

  gLastRpsTime = millis();

  Serial.println(F("──────────────────────────────"));
  Serial.println(F("   System Ready + IMU         "));
  Serial.println(F("──────────────────────────────"));
}

void loop() {
  // --- ToF-LIDAR (distance & motor PWM) ---
  uint8_t newData = 0;
  if (tof.VL53L4CX_GetMeasurementDataReady(&newData) != 0 || !newData) {
    // No new ToF data yet; IMU/fusion can still run!
    goto IMU_FUSION;
  }
  VL53L4CX_MultiRangingData_t data;
  if (tof.VL53L4CX_GetMultiRangingData(&data) != 0) {
    gConsecutiveReadErrors++;
#if TOF_DEBUG
    Serial.println(F("[ToF] Read error"));
#endif
#if USE_XSHUT_PIN
    if (gConsecutiveReadErrors >= kReadErrorThreshold) {
      sensorHardReset();
      gConsecutiveReadErrors = 0;
    }
#endif
    tof.VL53L4CX_ClearInterruptAndStartMeasurement();
    goto IMU_FUSION;
  }
  gConsecutiveReadErrors = 0;

  if (data.NumberOfObjectsFound == 0) {
    gNoTargetFrames++;
    gHasValidDistance = false;
    ledcWrite(kMotorPwmPin, 0);
    if (gNoTargetFrames > kNoTargetThreshold && gCurrentMode != MODE_LONG) {
      gCurrentMode = (gCurrentMode == MODE_SHORT) ? MODE_MEDIUM : MODE_LONG;
      applyMode(gCurrentMode);
      gNoTargetFrames = 0;
      tof.VL53L4CX_ClearInterruptAndStartMeasurement();
      goto IMU_FUSION;
    }
  }

  int bestIdx = -1;
  int16_t bestDelta = INT16_MAX;
  for (int i = 0; i < data.NumberOfObjectsFound; ++i) {
    const auto &r = data.RangeData[i];
    int16_t mm = r.RangeMilliMeter;
    if (gPrevDistanceMm < 0) {
      if (bestIdx < 0 || mm < data.RangeData[bestIdx].RangeMilliMeter) {
        bestIdx = i;
      }
    } else {
      int16_t delta = abs(mm - gPrevDistanceMm);
      if (delta < bestDelta) {
        bestDelta = delta;
        bestIdx   = i;
      }
    }
  }

  if (bestIdx >= 0) {
    const auto &best = data.RangeData[bestIdx];
    bool modeChanged = false;
    int16_t distMm = best.RangeMilliMeter;

    if (gCurrentMode == MODE_LONG && distMm < kMediumMaxMm) {
      bool wantShort = (distMm < kShortMaxMm);
      gStableTargetFrames++;
      if (gStableTargetFrames > kStableTargetThreshold) {
        gCurrentMode = wantShort ? MODE_SHORT : MODE_MEDIUM;
        applyMode(gCurrentMode);
        modeChanged = true;
        gStableTargetFrames = 0;
      }
    } else if (gCurrentMode == MODE_MEDIUM && distMm < kShortMaxMm) {
      gStableTargetFrames++;
      if (gStableTargetFrames > kStableTargetThreshold) {
        gCurrentMode = MODE_SHORT;
        applyMode(gCurrentMode);
        modeChanged = true;
        gStableTargetFrames = 0;
      }
    } else {
      gStableTargetFrames = 0;
    }

    gNoTargetFrames = 0;
    if (modeChanged) {
      tof.VL53L4CX_ClearInterruptAndStartMeasurement();
      goto IMU_FUSION;
    }

    if (best.RangeStatus == VL53L4CX_RANGESTATUS_RANGE_VALID ||
        best.RangeStatus == VL53L4CX_RANGESTATUS_RANGE_VALID_MIN_RANGE_CLIPPED ||
        best.RangeStatus == VL53L4CX_RANGESTATUS_RANGE_VALID_MERGED_PULSE)
    {
      Serial.print(static_cast<uint16_t>(best.RangeMilliMeter));
      Serial.print(F(" mm"));
      const int16_t min_dist_mm = 50;
      const int16_t max_dist_mm = 400;
      long pwm_value = map(distMm, min_dist_mm, max_dist_mm, 255, 0);
      pwm_value = constrain(pwm_value, 0, 255);
      ledcWrite(kMotorPwmPin, pwm_value);
      Serial.print(F(" -> PWM: "));
      Serial.print(pwm_value);
      float currentRPS = getRPS();
      Serial.print(F(" | RPS: "));
      Serial.print(currentRPS, 2);
      gPrevDistanceMm = best.RangeMilliMeter;
      gLatestDistanceMm = best.RangeMilliMeter;
      gHasValidDistance = true;
      Serial.print(F(" | "));
    } else {
      ledcWrite(kMotorPwmPin, 0);
      gHasValidDistance = false;
      Serial.print(F("-- mm -> PWM: 0 | RPS: -- | "));
    }
  }
  tof.VL53L4CX_ClearInterruptAndStartMeasurement();
  maybeReportMissingPi();

IMU_FUSION:
  // --- IMU Read + Fusion ---
  sensors_event_t accel, gyro, temp, mag;
  sox.getEvent(&accel, &gyro, &temp);
  lis3mdl.getEvent(&mag);

  // Convert gyro from rad/s to deg/s for fusion
  float gx = gyro.gyro.x * SENSORS_RADS_TO_DPS;
  float gy = gyro.gyro.y * SENSORS_RADS_TO_DPS;
  float gz = gyro.gyro.z * SENSORS_RADS_TO_DPS;

  filter.update(
    gx, gy, gz,
    accel.acceleration.x, accel.acceleration.y, accel.acceleration.z,
    mag.magnetic.x, mag.magnetic.y, mag.magnetic.z
  );

  gYaw = filter.getYaw();   if (gYaw < 0) gYaw += 360.0;
  gPitch = filter.getPitch(); if (gPitch < 0) gPitch += 360.0;
  gRoll = filter.getRoll();  if (gRoll < 0) gRoll += 360.0;

  Serial.print("Yaw: "); Serial.print(gYaw, 2); Serial.print(" deg, ");
  Serial.print("Pitch: "); Serial.print(gPitch, 2); Serial.print(" deg, ");
  Serial.print("Roll: "); Serial.print(gRoll, 2); Serial.println(" deg");

  delay(10); // ~100Hz update
}