#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <MAX30105.h>
#include "heartRate.h"

// ---------------- OLED SETTINGS -----------------
#define OLED_WIDTH 128
#define OLED_HEIGHT 64
#define OLED_ADDR 0x3C

Adafruit_SSD1306 display(OLED_WIDTH, OLED_HEIGHT, &Wire);

// ---------------- MAX30102 SETTINGS -------------
MAX30105 sensor;
long lastBeat = 0;
float bpm = 0;
bool beatFlag = false;

// ---------------- PIN SELECTION -----------------
#define SDA_PIN 5
#define SCL_PIN 6

// ---------------- SpO2 BUFFER -------------------
const int BUFFER_SIZE = 50;
long irBuf[BUFFER_SIZE];
long redBuf[BUFFER_SIZE];
int bufIndex = 0;
int bufCount = 0;
float spo2 = 0.0;
unsigned long lastSpO2Update = 0;

void computeSpO2() {
  if (bufCount < BUFFER_SIZE) return;

  double sumIR = 0, sumRED = 0;
  for (int i = 0; i < BUFFER_SIZE; i++) {
    sumIR += irBuf[i];
    sumRED += redBuf[i];
  }

  double dcIR = sumIR / BUFFER_SIZE;
  double dcRED = sumRED / BUFFER_SIZE;

  double sumSqIR = 0, sumSqRED = 0;
  for (int i = 0; i < BUFFER_SIZE; i++) {
    double acIR = irBuf[i] - dcIR;
    double acRED = redBuf[i] - dcRED;
    sumSqIR += acIR * acIR;
    sumSqRED += acRED * acRED;
  }

  double rmsIR = sqrt(sumSqIR / BUFFER_SIZE);
  double rmsRED = sqrt(sumSqRED / BUFFER_SIZE);

  if (dcIR <= 0 || dcRED <= 0 || rmsIR <= 0) return;

  double R = (rmsRED / dcRED) / (rmsIR / dcIR);
  double estSpO2 = 110.0 - 25.0 * R;

  if (estSpO2 < 70.0) estSpO2 = 70.0;
  if (estSpO2 > 100.0) estSpO2 = 100.0;

  spo2 = estSpO2;
}

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

  Wire.begin(SDA_PIN, SCL_PIN);
  Wire.setClock(400000);

  // --------- MAX30102 FIRST ----------
  Serial.println("Initializing MAX30102...");
  if (!sensor.begin(Wire, I2C_SPEED_FAST)) {
    Serial.println("MAX30102 not found");
    while (1);
  }

  sensor.setup();
  sensor.setLEDMode(2);              // RED + IR
  sensor.setADCRange(16384);
  sensor.setSampleRate(100);
  sensor.setPulseWidth(411);

  // ⭐ MEDIUM LED brightness (best for waveform)
  sensor.setPulseAmplitudeRed(0x3F);
  sensor.setPulseAmplitudeIR(0x3F);

  // --------- OLED SECOND ------------
  Serial.println("Initializing OLED...");
  if (!display.begin(SSD1306_SWITCHCAPVCC, OLED_ADDR)) {
    Serial.println("OLED not found");
    while (1);
  }

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.display();

  Serial.println("System ready!");
}

void loop() {
  long ir  = sensor.getIR();
  long red = sensor.getRed();

  // ---------- HEART RATE CALC ----------
  beatFlag = false;
  if (checkForBeat(ir)) {
    long delta = millis() - lastBeat;
    lastBeat = millis();
    bpm = 60.0 / (delta / 1000.0);
    beatFlag = true;
  }

  // ---------- SpO2 BUFFER FILL ----------
  irBuf[bufIndex]  = ir;
  redBuf[bufIndex] = red;
  bufIndex = (bufIndex + 1) % BUFFER_SIZE;
  if (bufCount < BUFFER_SIZE) bufCount++;

  if (millis() - lastSpO2Update > 300) {
    computeSpO2();
    lastSpO2Update = millis();
  }

  // ---------- SERIAL PLOTTER ----------
  // Format: IR, RED, BPM
  // This makes 3 beautiful lines in Serial Plotter
  Serial.print(ir);
  Serial.print(",");
  Serial.print(red);
  Serial.print(",");
  Serial.println(bpm);

  // ---------- OLED OUTPUT ----------
  display.clearDisplay();

  display.setTextSize(1);
  display.setCursor(0, 0);
  display.print("IR: ");
  display.println(ir);

  display.setCursor(0, 12);
  display.print("RED: ");
  display.println(red);

  display.setTextSize(2);
  display.setCursor(0, 28);
  display.print("BPM:");
  display.println((int)bpm);

  display.setTextSize(1);
  display.setCursor(0, 52);
  display.print("SpO2: ");
  display.print(spo2, 1);
  display.print("%");

  if (beatFlag) {
    display.fillCircle(110, 10, 3, SSD1306_WHITE);
  }

  display.display();

  delay(20);
}