#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include "esp_camera.h"
#include "facettd_inferencing.h"
#include "edge-impulse-sdk/classifier/ei_run_classifier.h"

// --- Prevent sensor_t conflict with ESP32 camera ---
#define sensor_t adafruit_sensor_t
#include <Adafruit_MPU6050.h>
#include <Adafruit_Sensor.h>
#undef sensor_t
#include "HX711.h"

// ===== Camera model & pins =====
#define CAMERA_MODEL_XIAO_ESP32S3
#include "camera_pins.h"

// ===== OLED =====
#define SCREEN_WIDTH   128
#define SCREEN_HEIGHT  64
#define SCREEN_ADDRESS 0x3C
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// ===== Button / LED =====
#define BUTTON_PIN  4
#define LED_PIN     3

// ===== I2C & HX711 Pins =====
#define SDA_PIN      5
#define SCL_PIN      6
#define HX711_DT     A8
#define HX711_SCK    A9

// ===== Modes =====
enum Mode { MENU, CAMERA, SENSORS };
Mode currentMode = MENU;

// ===== Button handling =====
unsigned long pressStart = 0;
unsigned long lastPressTime = 0;
int pressCount = 0;
bool buttonHeld = false;
const unsigned long PRESS_TIMEOUT = 1000;
const unsigned long HOLD_THRESHOLD = 800;

// ===== Camera buffers =====
static float gray_buffer[SCREEN_WIDTH * SCREEN_HEIGHT];
static uint8_t rgb888_buf[160 * 120 * 3];
static float ei_input_buf[
  EI_CLASSIFIER_INPUT_WIDTH *
  EI_CLASSIFIER_INPUT_HEIGHT *
  EI_CLASSIFIER_INPUT_FRAMES
];
bool camera_ready = false;

// ===== Sensors =====
Adafruit_MPU6050 mpu;
HX711 scale;
//float calibration_factor = 2280.0f;
float calibration_factor = 1.0f;

// ---------- Text Helper ----------
void drawTextCentered(const char* line1, const char* line2 = nullptr) {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  int16_t x1, y1; uint16_t w, h;
  if (line1) {
    display.getTextBounds(line1, 0, 0, &x1, &y1, &w, &h);
    display.setCursor((SCREEN_WIDTH - w) / 2, (SCREEN_HEIGHT / 2) - (line2 ? 12 : h / 2));
    display.println(line1);
  }
  if (line2) {
    display.getTextBounds(line2, 0, 0, &x1, &y1, &w, &h);
    display.setCursor((SCREEN_WIDTH - w) / 2, (SCREEN_HEIGHT / 2) + 2);
    display.println(line2);
  }
  display.display();
}

// ---------- Menu ----------
void showMenu() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);

  // Move "Press:" to the very top
  display.setCursor(5, 0);
  display.println("Press:");

  // Space options evenly for readability
  display.setCursor(10, 14);
  display.println("1x -> Camera");
  display.setCursor(10, 26);
  display.println("2x -> Sensors");
  display.setCursor(10, 38);
  display.println("3x -> Menu");

  // Keep flashlight hint at the bottom edge
  display.setCursor(10, 54);
  display.println("Hold -> Flashlight");

  display.display();
}

// ---------- Edge Impulse ----------
void prepare_ei_input_from_rgb888(int sw, int sh) {
  const float sx = (float)sw / EI_CLASSIFIER_INPUT_WIDTH;
  const float sy = (float)sh / EI_CLASSIFIER_INPUT_HEIGHT;
  int dst = 0;
  for (int y = 0; y < EI_CLASSIFIER_INPUT_HEIGHT; y++) {
    int yy = (int)(y * sy);
    for (int x = 0; x < EI_CLASSIFIER_INPUT_WIDTH; x++) {
      int xx = (int)(x * sx);
      const uint8_t* p = rgb888_buf + (yy * sw + xx) * 3;
      ei_input_buf[dst++] = (float)((p[0] << 16) | (p[1] << 8) | p[2]);
    }
  }
}

float run_inference() {
  signal_t signal;
  signal.total_length = sizeof(ei_input_buf) / sizeof(float);
  signal.get_data = [](size_t offset, size_t length, float* out_ptr)->int {
    memcpy(out_ptr, ei_input_buf + offset, length * sizeof(float));
    return 0;
  };

  ei_impulse_result_t result = {0};
  EI_IMPULSE_ERROR r = run_classifier(&signal, &result, false);
  if (r != EI_IMPULSE_OK) {
    Serial.printf("run_classifier failed (%d)\n", r);
    return -1;
  }
  float val = result.classification[0].value;
  Serial.printf("TTD: %.3f years\n", val);
  return val;
}

// ---------- Dithered Display ----------
static void processAndDisplayImage(const camera_fb_t* fb) {
  const uint8_t* src = fb->buf;
  const int sw = fb->width, sh = fb->height;
  const float sx = (float)sw / SCREEN_WIDTH;
  const float sy = (float)sh / SCREEN_HEIGHT;

  for (int y = 0; y < SCREEN_HEIGHT; ++y) {
    int y0 = (int)(y * sy), y1 = (int)((y + 1) * sy);
    if (y1 <= y0) y1 = y0 + 1;
    for (int x = 0; x < SCREEN_WIDTH; ++x) {
      int x0 = (int)(x * sx), x1 = (int)((x + 1) * sx);
      if (x1 <= x0) x1 = x0 + 1;
      uint32_t sum = 0; int cnt = 0;
      for (int yy = y0; yy < y1; yy++) {
        const int row = yy * sw;
        for (int xx = x0; xx < x1; xx++) { sum += src[row + xx]; cnt++; }
      }
      gray_buffer[y * SCREEN_WIDTH + x] = (cnt ? (float)sum / cnt : 0.0f);
    }
  }

  float gmin = 255, gmax = 0;
  for (int i = 0; i < SCREEN_WIDTH * SCREEN_HEIGHT; i++) {
    float v = gray_buffer[i];
    if (v < gmin) gmin = v;
    if (v > gmax) gmax = v;
  }
  float range = gmax - gmin; if (range < 1) range = 1;
  for (int i = 0; i < SCREEN_WIDTH * SCREEN_HEIGHT; i++)
    gray_buffer[i] = (gray_buffer[i] - gmin) * (255.0f / range);

  // Floyd–Steinberg dithering
  for (int y = 0; y < SCREEN_HEIGHT; y++) {
    for (int x = 0; x < SCREEN_WIDTH; x++) {
      int idx = y * SCREEN_WIDTH + x;
      float oldp = gray_buffer[idx];
      float newp = oldp >= 128 ? 255 : 0;
      float err = oldp - newp;
      gray_buffer[idx] = newp;
      if (x + 1 < SCREEN_WIDTH) gray_buffer[idx + 1] += err * 7 / 16.0f;
      if (y + 1 < SCREEN_HEIGHT) {
        if (x > 0) gray_buffer[idx + SCREEN_WIDTH - 1] += err * 3 / 16.0f;
        gray_buffer[idx + SCREEN_WIDTH] += err * 5 / 16.0f;
        if (x + 1 < SCREEN_WIDTH) gray_buffer[idx + SCREEN_WIDTH + 1] += err * 1 / 16.0f;
      }
    }
  }

  display.clearDisplay();
  for (int y = 0; y < SCREEN_HEIGHT; y++)
    for (int x = 0; x < SCREEN_WIDTH; x++)
      if (gray_buffer[y * SCREEN_WIDTH + x] > 0.5f)
        display.drawPixel(x, y, SSD1306_WHITE);
  display.display();
}

// ---------- Overlay Result ----------
void drawTTDBox(float years) {
  display.fillRect(0, SCREEN_HEIGHT - 12, SCREEN_WIDTH, 12, SSD1306_BLACK);
  display.drawRect(0, SCREEN_HEIGHT - 12, SCREEN_WIDTH, 12, SSD1306_WHITE);
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(4, SCREEN_HEIGHT - 10);
  display.printf("TTD: %.2f years", years);
  display.display();
  delay(2000);  // keep result visible for 2 seconds
}

// ---------- CAMERA MODE ----------
void runCameraMode() {
  drawTextCentered("Camera Mode", "Press 3x to return");

  // --- Camera config ---
  camera_config_t config;
  config.ledc_channel = LEDC_CHANNEL_0;
  config.ledc_timer = LEDC_TIMER_0;
  config.pin_d0 = Y2_GPIO_NUM; config.pin_d1 = Y3_GPIO_NUM;
  config.pin_d2 = Y4_GPIO_NUM; config.pin_d3 = Y5_GPIO_NUM;
  config.pin_d4 = Y6_GPIO_NUM; config.pin_d5 = Y7_GPIO_NUM;
  config.pin_d6 = Y8_GPIO_NUM; config.pin_d7 = Y9_GPIO_NUM;
  config.pin_xclk = XCLK_GPIO_NUM; config.pin_pclk = PCLK_GPIO_NUM;
  config.pin_vsync = VSYNC_GPIO_NUM; config.pin_href = HREF_GPIO_NUM;
  config.pin_sscb_sda = SIOD_GPIO_NUM; config.pin_sscb_scl = SIOC_GPIO_NUM;
  config.pin_pwdn = PWDN_GPIO_NUM; config.pin_reset = RESET_GPIO_NUM;
  config.xclk_freq_hz = 20000000;
  config.frame_size = FRAMESIZE_QQVGA;
  config.pixel_format = PIXFORMAT_GRAYSCALE;
  config.fb_location = CAMERA_FB_IN_PSRAM;
  config.jpeg_quality = 12;
  config.fb_count = 1;
  config.grab_mode = CAMERA_GRAB_LATEST;

  if (esp_camera_init(&config) == ESP_OK) {
    camera_ready = true;
    auto* s = esp_camera_sensor_get();
    s->set_vflip(s, 1);
    s->set_hmirror(s, 0);
  } else {
    drawTextCentered("Camera init failed");
    delay(1500);
    currentMode = MENU;
    return;
  }

  while (currentMode == CAMERA) {
    bool pressed = (digitalRead(BUTTON_PIN) == LOW);
    unsigned long now = millis();
    if (pressed && pressStart == 0) pressStart = now;
    if (pressed && !buttonHeld && (now - pressStart > HOLD_THRESHOLD)) {
      buttonHeld = true; digitalWrite(LED_PIN, HIGH);
    }
    if (!pressed && pressStart > 0) {
      if (buttonHeld) { digitalWrite(LED_PIN, LOW); buttonHeld = false; }
      else { pressCount++; lastPressTime = now; }
      pressStart = 0;
    }

    if (pressCount > 0 && (now - lastPressTime > PRESS_TIMEOUT)) {
      if (pressCount == 1) {
        camera_fb_t* fb_inf = esp_camera_fb_get();
        if (fb_inf) {
          fmt2rgb888(fb_inf->buf, fb_inf->len, PIXFORMAT_GRAYSCALE, rgb888_buf);
          prepare_ei_input_from_rgb888(fb_inf->width, fb_inf->height);
          float ttd = run_inference();
          drawTTDBox(ttd);
          esp_camera_fb_return(fb_inf);
        }
      } else if (pressCount == 2) currentMode = SENSORS;
      else if (pressCount >= 3) currentMode = MENU;
      pressCount = 0;
    }

    camera_fb_t* fb = esp_camera_fb_get();
    if (fb) {
      processAndDisplayImage(fb);
      esp_camera_fb_return(fb);
    }
    delay(30);
  }
  esp_camera_deinit();
}

// ---------- SENSOR MODE ----------
void runSensorMode() {
  drawTextCentered("Sensors Mode", "Press 3x to return");
  Wire.begin(SDA_PIN, SCL_PIN);
  scale.begin(HX711_DT, HX711_SCK);
  scale.set_scale(calibration_factor);
  scale.tare();

  unsigned long lastRead = 0;
  bool collecting = false;
  bool mpu_ready = false;

  while (currentMode == SENSORS) {
    bool pressed = (digitalRead(BUTTON_PIN) == LOW);
    unsigned long now = millis();
    if (pressed && pressStart == 0) pressStart = now;
    if (pressed && !buttonHeld && (now - pressStart > HOLD_THRESHOLD)) {
      buttonHeld = true; digitalWrite(LED_PIN, HIGH);
    }
    if (!pressed && pressStart > 0) {
      if (buttonHeld) { digitalWrite(LED_PIN, LOW); buttonHeld = false; }
      else { pressCount++; lastPressTime = now; }
      pressStart = 0;
    }
    if (pressCount > 0 && (now - lastPressTime > PRESS_TIMEOUT)) {
      if (pressCount >= 3) currentMode = MENU;
      pressCount = 0;
    }

    // Begin data collection when button held down
    if (pressed) {
      if (!collecting) {
        collecting = true;
        Serial.println("=== START DATA COLLECTION ===");

        // --- Try discovering MPU6050 again ---
        mpu_ready = false;
        for (uint8_t addr : {0x68, 0x69}) {
          if (mpu.begin(addr)) {
            Serial.printf("MPU6050 found at 0x%02X\n", addr);
            mpu_ready = true;
            break;
          }
        }
        if (!mpu_ready) {
          Serial.println("MPU6050 not detected! Check wiring.");
          display.clearDisplay();
          display.setCursor(0, 0);
          display.println("MPU6050 not found!");
          display.display();
        }
      }

      if (now - lastRead >= 500) {
        lastRead = now;
        sensors_event_t a, g, temp;
        bool ok = mpu_ready && mpu.getEvent(&a, &g, &temp);

        long raw = 0; float weight = 0;
        if (scale.is_ready()) { raw = scale.read(); weight = scale.get_units(5); }

        display.clearDisplay();
        display.setTextSize(1);
        display.setCursor(0, 0);
        if (ok) {
          display.printf("Ax: %.1f Ay: %.1f\n", a.acceleration.x, a.acceleration.y);
          display.printf("Az: %.1f T: %.1fC\n", a.acceleration.z, temp.temperature);
          display.printf("Gx: %.1f Gy: %.1f\n", g.gyro.x, g.gyro.y);
          display.printf("W: %.1f g\n", weight);
        } else {
          display.println("MPU read fail");
        }
        display.display();
      }
    }
    else if (collecting) {
      collecting = false;
      Serial.println("=== STOP DATA COLLECTION ===");
      display.clearDisplay();
      display.setCursor(0, 0);
      display.println("Data collection stopped.");
      display.display();
    }

    delay(10);
  }
}

// ---------- SETUP ----------
void setup() {
  Serial.begin(115200);
  pinMode(LED_PIN, OUTPUT);
  pinMode(BUTTON_PIN, INPUT_PULLUP);
  Wire.begin(SDA_PIN, SCL_PIN);
  if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS))
    Serial.println("OLED init failed!");
  showMenu();
}

// ---------- LOOP ----------
void loop() {
  bool pressed = (digitalRead(BUTTON_PIN) == LOW);
  unsigned long now = millis();
  if (pressed && pressStart == 0) pressStart = now;
  if (pressed && !buttonHeld && (now - pressStart > HOLD_THRESHOLD)) {
    buttonHeld = true; digitalWrite(LED_PIN, HIGH);
  }
  if (!pressed && pressStart > 0) {
    if (buttonHeld) { digitalWrite(LED_PIN, LOW); buttonHeld = false; }
    else { pressCount++; lastPressTime = now; }
    pressStart = 0;
  }
  if (pressCount > 0 && (now - lastPressTime > PRESS_TIMEOUT)) {
    if (pressCount == 1) currentMode = CAMERA;
    else if (pressCount == 2) currentMode = SENSORS;
    else if (pressCount >= 3) currentMode = MENU;
    pressCount = 0;
  }

  switch (currentMode) {
    case MENU:    showMenu(); break;
    case CAMERA:  runCameraMode(); break;
    case SENSORS: runSensorMode(); break;
  }
}