#include <Wire.h>
#include <Adafruit_SSD1306.h>

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET    -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

const int ecgPin = 4;  // GPIO4 where the ECG signal is input

// LED and Buzzer pins
const int ledLowPin = 2;     // LED for low heart rate (GPIO2)
const int ledNormalPin = 3;  // LED for normal heart rate (GPIO3)
const int ledHighPin = 0;    // LED for high heart rate (GPIO0)
const int buzzerPin = 10;     // Buzzer pin (GPIO8)

// Thresholds for heart rate (adjust based on your needs)
const int lowThreshold = 120;     // Below 120 BPM is LOW
const int highThreshold = 187;   // Above 187 BPM is HIGH

void setup() {
  Serial.begin(115200);
  pinMode(ecgPin, INPUT);
  
  // Initialize LED and Buzzer pins
  pinMode(ledLowPin, OUTPUT);
  pinMode(ledNormalPin, OUTPUT);
  pinMode(ledHighPin, OUTPUT);
  pinMode(buzzerPin, OUTPUT);
  
  // Start with all LEDs off
  digitalWrite(ledLowPin, LOW);
  digitalWrite(ledNormalPin, LOW);
  digitalWrite(ledHighPin, LOW);
  digitalWrite(buzzerPin, LOW);

  // Initialize the OLED display
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);
  }
  
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
}

// Function to calculate BPM from ECG samples (simplified)
int calculateBPM(int *ecgValues, int length) {
  int peaks = 0;
  int lastPeak = -10; // Avoid counting same peak multiple times
  
  // Simple peak detection
  for (int i = 2; i < length - 2; i++) {
    if (ecgValues[i] > ecgValues[i-1] && 
        ecgValues[i] > ecgValues[i-2] &&
        ecgValues[i] > ecgValues[i+1] && 
        ecgValues[i] > ecgValues[i+2] &&
        ecgValues[i] > SCREEN_HEIGHT / 2) { // Peak threshold
      
      if (i - lastPeak > 10) { // Minimum distance between peaks
        peaks++;
        lastPeak = i;
      }
    }
  }
  
  // Calculate BPM based on peaks in the sample window
  // Assuming the screen width represents about 10 seconds of data
  // (adjust based on your delay timing: 128 samples * 10ms = 1.28 seconds)
  // For more accurate BPM, adjust the time factor
  int bpm = peaks * 60 / 1.28; // Convert peaks per 1.28 seconds to BPM
  
  return constrain(bpm, 30, 200); // Constrain to reasonable values
}

// Function to update LEDs and buzzer based on BPM
void updateIndicators(int bpm) {
  // Turn off all LEDs first
  digitalWrite(ledLowPin, LOW);
  digitalWrite(ledNormalPin, LOW);
  digitalWrite(ledHighPin, LOW);
  
  // Determine heart rate category and activate appropriate LED
  if (bpm < lowThreshold) {
    // LOW heart rate
    digitalWrite(ledLowPin, HIGH);
    // Buzzer pattern: slow beep (1 beep per second)
    static unsigned long lastBeep = 0;
    if (millis() - lastBeep > 1000) {
      digitalWrite(buzzerPin, HIGH);
      delay(100);
      digitalWrite(buzzerPin, LOW);
      lastBeep = millis();
    }
    Serial.print("LOW Heart Rate: ");
  } 
  else if (bpm > highThreshold) {
    // HIGH heart rate
    digitalWrite(ledHighPin, HIGH);
    // Buzzer pattern: fast beep (2 beeps per second)
    static unsigned long lastBeep = 0;
    if (millis() - lastBeep > 500) {
      digitalWrite(buzzerPin, HIGH);
      delay(100);
      digitalWrite(buzzerPin, LOW);
      lastBeep = millis();
    }
    Serial.print("HIGH Heart Rate: ");
  } 
  else {
    // NORMAL heart rate
    digitalWrite(ledNormalPin, HIGH);
    // Buzzer: short acknowledgment beep every 5 seconds
    static unsigned long lastBeep = 0;
    if (millis() - lastBeep > 5000) {
      digitalWrite(buzzerPin, HIGH);
      delay(50);
      digitalWrite(buzzerPin, LOW);
      lastBeep = millis();
    }
    Serial.print("NORMAL Heart Rate: ");
  }
  
  Serial.print(bpm);
  Serial.println(" BPM");
}

void loop() {
  // Array to hold ECG samples for displaying
  int ecgValues[SCREEN_WIDTH];

  // Read ECG signal from custom chip and store values
  for (int i = 0; i < SCREEN_WIDTH; i++) {
    int ecgSignal = analogRead(ecgPin);  // Capture ECG signal from GPIO4
    ecgValues[i] = map(ecgSignal, 0, 4095, 0, SCREEN_HEIGHT);  // Scale to screen height
    delay(10);  // Small delay for smoother waveform (tune this based on your signal speed)
  }

  // Calculate BPM from the collected ECG samples
  int bpm = calculateBPM(ecgValues, SCREEN_WIDTH);
  
  // Update LEDs and buzzer based on calculated BPM
  updateIndicators(bpm);

  // Clear the display
  display.clearDisplay();

  // Draw ECG waveform
  for (int i = 1; i < SCREEN_WIDTH; i++) {
    display.drawLine(i - 1, ecgValues[i - 1], i, ecgValues[i], SSD1306_WHITE);
  }

  // Display BPM and status on OLED
  display.setCursor(0, 0);
  display.print("ECG Signal - ");
  
  // Show status on display
  if (bpm < lowThreshold) {
    display.print("LOW");
  } else if (bpm > highThreshold) {
    display.print("HIGH");
  } else {
    display.print("NORMAL");
  }
  
  display.setCursor(0, 8);
  display.print("BPM: ");
  display.print(bpm);

  // Update the display with new data
  display.display();

  delay(10);  // Small delay before next loop
}