#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include "esp_camera.h"

// ===== Edge Impulse =====
#include "facettd_inferencing.h"
#include "edge-impulse-sdk/classifier/ei_run_classifier.h"

// ===== Camera model & pins (XIAO ESP32S3 Sense) =====
#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);

// ===== LED & Touch =====
#define PIN_LED 21
#define N_TOUCH 2
#define THRESHOLD 100000UL
int touch_pins[N_TOUCH] = {1, 2};
bool pin_touched_now[N_TOUCH];
bool pin_touched_past[N_TOUCH];
bool justPressed(int i){ return pin_touched_now[i] && !pin_touched_past[i]; }

// ===== Timing =====
const unsigned long CAPTURE_INTERVAL_MS = 60000UL;
unsigned long lastCapture = 0;
bool camera_ready = false;

// ===== Buffers =====
static float gray_buffer[SCREEN_WIDTH * SCREEN_HEIGHT];
static uint8_t rgb888_buf[160 * 120 * 3]; // for inference
static float ei_input_buf[
  EI_CLASSIFIER_INPUT_WIDTH *
  EI_CLASSIFIER_INPUT_HEIGHT *
  EI_CLASSIFIER_INPUT_FRAMES
];

// ---------- Touch Update ----------
void update_touch(){
  for(int i=0;i<N_TOUCH;i++){
    uint32_t v=touchRead(touch_pins[i]);
    pin_touched_past[i]=pin_touched_now[i];
    pin_touched_now[i]=(v>THRESHOLD);
  }
}

// ---------- Text Helper ----------
static 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();
}

// ---------- Edge Impulse helpers ----------
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;
}

// ---------- Original Grayscale + Dithering 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 dither
  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();
}

// ---------- Setup ----------
void setup(){
  Serial.begin(115200);
  delay(250);
  pinMode(PIN_LED,OUTPUT);
  digitalWrite(PIN_LED,HIGH);
  Wire.begin();
  delay(50);

  if(!display.begin(SSD1306_SWITCHCAPVCC,SCREEN_ADDRESS))
    Serial.println("SSD1306 init failed");
  else
    drawTextCentered("Camera + Touch","Init...");

  // Camera config: dual path (gray for OLED, RGB565 for inference)
  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;
    sensor_t* s=esp_camera_sensor_get();
    s->set_vflip(s,1);
    s->set_hmirror(s,0);
    Serial.println("Camera ready");
    drawTextCentered("Camera Ready","Touch or Wait 60s");
    lastCapture=millis();
  }else{
    drawTextCentered("Camera init failed");
  }
}

// ---------- Loop ----------
void loop(){
  update_touch();
  bool touch_trigger=justPressed(0)||justPressed(1);
  if(!camera_ready){delay(200);return;}

  unsigned long now=millis();
  if((now-lastCapture>=CAPTURE_INTERVAL_MS)||touch_trigger){
    lastCapture=now;
    digitalWrite(PIN_LED,LOW);
    drawTextCentered("Capturing...");
    camera_fb_t* fb=esp_camera_fb_get();
    digitalWrite(PIN_LED,HIGH);

    if(!fb){
      Serial.println("Capture failed");
      drawTextCentered("Capture failed");
      return;
    }

    Serial.printf("Captured %dx%d (%d bytes)\n",fb->width,fb->height,fb->len);

    // 1️⃣ Show perfect grayscale dithering
    processAndDisplayImage(fb);

    // 2️⃣ Convert to RGB888 for inference (parallel)
    fmt2rgb888(fb->buf,fb->len,PIXFORMAT_GRAYSCALE,rgb888_buf);
    prepare_ei_input_from_rgb888(fb->width,fb->height);
    float ttd_years=run_inference();

    // 3️⃣ Overlay result
    drawTTDBox(ttd_years);

    esp_camera_fb_return(fb);
  }
  delay(30);
}