#include <Wire.h>
#include <EEPROM.h>
#include "MAX30105.h"
#include <BLEDevice.h>
#include <BLEServer.h>
#include <BLEUtils.h>
#include <BLE2902.h>
#include <ESP32Servo.h>
#include <U8g2lib.h>
#include <pgmspace.h>

// — I²C BUS PINS —  
#define I2C_SDA 5  
#define I2C_SCL 6  

// — visible panel is 72×40 centered in 128×64 RAM —  
const int OFF_X = 28;   // (128−72)/2  
const int OFF_Y = 21;   // (64−40)/2  

// — OLED (SSD1306 128×64) via U8g2 —  
U8G2_SSD1306_128X64_NONAME_F_HW_I2C u8g2(
  U8G2_R0, U8X8_PIN_NONE, I2C_SCL, I2C_SDA
);

// — 16×16 px Bluetooth icon (row-major, LSB-first) —  
#define BT_W 16
#define BT_H 16
static const uint8_t bluetooth16_bits[] PROGMEM = {
0x80, 0x01, 0x80, 0x02, 0x80, 0x0c, 0x80, 0x10, 0x88, 0x20, 0x90, 0x10, 0xa0, 0x0c, 0xc0, 0x02, 
0xc0, 0x01, 0xa0, 0x02, 0x90, 0x0c, 0x88, 0x10, 0x80, 0x20, 0x80, 0x18, 0x80, 0x04, 0x80, 0x03
};

// — PROTOTYPES —  
int  runInference(int sensorVal);
int  argMax(float* arr,int length);
void controlMotor(int c);
void stopMotor();

// — BUTTON & DEBOUNCE —  
#define BUTTON_PIN     21  
bool  bleEnabled      = true;  
bool  lastButtonState = HIGH;  
bool  buttonState     = HIGH;  
unsigned long lastDebounceTime = 0;  
const unsigned long debounceDelay = 50;

// — SERVO (SG90) —  
Servo myServo;  
const int servoPin    = 4;  
int   currentAngle    = 90;  
const int stepDelayMs = 5;

// — MAX30105 SENSOR —  
MAX30105 particleSensor;  
#define SCALE_FACTOR 100  
bool sensorFound = false;

// — MODEL PARAMETERS —  
float parameters[16];  
bool  paramsStored = false;

// — BLE GLOBALS —  
BLEServer*         pServer         = nullptr;  
BLECharacteristic* pCharacteristic = nullptr;  
BLEAdvertising*    pAdvertising    = nullptr;  
bool               deviceConnected = false;

// ===== BLE CALLBACKS =====  
class MyServerCallbacks : public BLEServerCallbacks {
  void onConnect(BLEServer*)  override { deviceConnected = true;  }
  void onDisconnect(BLEServer*) override {
    deviceConnected = false;
    if (bleEnabled) pAdvertising->start();
  }
};

class ParamCallbacks : public BLECharacteristicCallbacks {
  void onWrite(BLECharacteristic* chr) override {
    String s = chr->getValue();
    float arr[16]; int cnt=0;
    char buf[s.length()+1];
    s.toCharArray(buf,sizeof(buf));
    char* tok = strtok(buf,",");
    while(tok && cnt<16){
      arr[cnt++]=atof(tok);
      tok=strtok(nullptr,",");
    }
    if(cnt==16){
      for(int i=0;i<16;i++) parameters[i]=arr[i];
      EEPROM.put(0,parameters);
      EEPROM.commit();
      paramsStored=true;
      // confirm
      u8g2.clearBuffer();
      u8g2.drawStr(OFF_X, OFF_Y+20, "Params Saved!");
      u8g2.sendBuffer();
      delay(1000);
    }
  }
};

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

  // 1) I²C @ 100 kHz  
  Wire.begin(I2C_SDA,I2C_SCL);
  Wire.setClock(100000);

  // 2) SSD1306 startup delay  
  delay(200);

  // 3) init OLED  
  u8g2.begin();
  u8g2.setFont(u8g2_font_6x10_tr);
  u8g2.clearBuffer();
  u8g2.sendBuffer();

  // 4) init servo  
  myServo.setPeriodHertz(50);
  myServo.attach(servoPin,500,2500);
  myServo.write(currentAngle);

  // 5) button  
  pinMode(BUTTON_PIN, INPUT_PULLUP);

  // 6) init sensor  
  sensorFound = particleSensor.begin(Wire,I2C_SPEED_STANDARD)
             || particleSensor.begin(Wire,I2C_SPEED_FAST);
  if(sensorFound){
    particleSensor.setup(0x1F,4,2,50,411,4096);
    particleSensor.setPulseAmplitudeRed(0);
  } else {
    Serial.println("⚠️ Sensor not found");
  }

  // 7) load model params  
  EEPROM.get(0,parameters);
  if(parameters[15]!=0.0f) paramsStored=true;

  // 8) init BLE  
  BLEDevice::init("Your Bionic Hand");
  pServer = BLEDevice::createServer();
  pServer->setCallbacks(new MyServerCallbacks());
  auto svc = pServer->createService("4fafc201-1fb5-459e-8fcc-c5c9c331914b");
  pCharacteristic = svc->createCharacteristic(
    "beb5483e-36e1-4688-b7f5-ea07361b26a8",
    BLECharacteristic::PROPERTY_READ   |
    BLECharacteristic::PROPERTY_WRITE  |
    BLECharacteristic::PROPERTY_NOTIFY
  );
  pCharacteristic->addDescriptor(new BLE2902());
  pCharacteristic->setCallbacks(new ParamCallbacks());
  pCharacteristic->setValue("Waiting...");
  svc->start();
  pAdvertising = pServer->getAdvertising();
  pAdvertising->addServiceUUID("4fafc201-1fb5-459e-8fcc-c5c9c331914b");
  pAdvertising->start();
  Serial.println("BLE advertising started");
}

void loop(){
  // — debounce button → toggle BLE —
  bool reading = digitalRead(BUTTON_PIN);
  if(reading!=lastButtonState) lastDebounceTime=millis();
  if(millis()-lastDebounceTime>debounceDelay && reading!=buttonState){
    buttonState=reading;
    if(buttonState==HIGH){
      bleEnabled = !bleEnabled;
      if(bleEnabled) pAdvertising->start();
      else           pAdvertising->stop();
    }
  }
  lastButtonState=reading;

  // — read sensor & inference —
  long rawIR = sensorFound ? particleSensor.getIR() : 0;
  int  scaled = rawIR / SCALE_FACTOR;
  int  pred   = paramsStored ? runInference(scaled) : -1;
  if(pred>=0) controlMotor(pred);
  else        stopMotor();

  // — update OLED (icon + status) —
  u8g2.clearBuffer();
  u8g2.setCursor(OFF_X, OFF_Y+10);
  u8g2.print("Sensor:"); u8g2.print(scaled);
  u8g2.setCursor(OFF_X, OFF_Y+20);
  u8g2.print("Predict:");u8g2.print(pred<0?-1:pred);

  // draw 16×16 Bluetooth icon, then ON/OFF
  u8g2.drawXBMP(OFF_X+18,   OFF_Y+24, BT_W, BT_H, bluetooth16_bits);
  u8g2.setCursor(OFF_X+36, OFF_Y+36);
  u8g2.print(bleEnabled ? "ON" : "OFF");

  u8g2.sendBuffer();

  // — BLE notify —
  if(bleEnabled && deviceConnected){
    char buf[8];
    int n = snprintf(buf,sizeof(buf),"%d",scaled);
    pCharacteristic->setValue((uint8_t*)buf,n);
    pCharacteristic->notify();
  }

  delay(100);
}

// ===== inference & servo ctrl (unchanged) =====
int argMax(float* a,int L){int idx=0;for(int i=1;i<L;i++)if(a[i]>a[idx])idx=i;return idx;}
int runInference(int v){
  float mean=parameters[14],stddev=parameters[15],x=(v-mean)/stddev,h[3],logits[2];
  for(int i=0;i<3;i++){ float z=parameters[i]*x+parameters[3+i]; h[i]=z>0?z:0; }
  for(int o=0;o<2;o++){ float sum=0; for(int i=0;i<3;i++) sum+=h[i]*parameters[6+i*2+o]; logits[o]=sum+parameters[12+o]; }
  return argMax(logits,2);
}
void controlMotor(int c){ int target=(c==0)?90:180,dir=(target>currentAngle)?1:-1;
  while(currentAngle!=target){ currentAngle+=dir; myServo.write(currentAngle); delay(stepDelayMs); }
}
void stopMotor(){ controlMotor(0); }