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

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define GRAPH_WIDTH 128
#define GRAPH_HEIGHT 40
#define GRAPH_Y_OFFSET 20

Adafruit_ADS1115 ads;
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

float dataPoints[GRAPH_WIDTH];
int currentX = 0;
float visualGain = 10.0;

unsigned long lastSample = 0;
unsigned long sampleCount = 0;

void setup() {
  Serial.begin(115200);
  Wire.begin();

  if (!ads.begin()) {
    Serial.println("ADS1115 error!");
    while (1);
  }

  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println("Display error!");
    while (1);
  }

  ads.setGain(GAIN_ONE);
  ads.setDataRate(RATE_ADS1115_860SPS);

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.println("EEG ready");
  display.display();

  Serial.println("timestamp,raw,voltage");
}

void loop() {
  unsigned long now = micros();

  if (now - lastSample >= 4000) {
    lastSample = now;
    sampleCount++;

    unsigned long timestamp = millis();
    int16_t raw = ads.readADC_SingleEnded(0);
    float voltage = raw * 4.096 / 32767.0;

    Serial.print(timestamp);
    Serial.print(",");
    Serial.print(raw);
    Serial.print(",");
    Serial.println(voltage, 6);

    float amplified = voltage * visualGain;

    dataPoints[currentX] = amplified;
    currentX = (currentX + 1) % GRAPH_WIDTH;

    drawWaveform();
  }
}

void drawWaveform() {
  display.clearDisplay();

  float minVal = dataPoints[0];
  float maxVal = dataPoints[0];
  for (int i = 1; i < GRAPH_WIDTH; i++) {
    if (dataPoints[i] < minVal) minVal = dataPoints[i];
    if (dataPoints[i] > maxVal) maxVal = dataPoints[i];
  }

  float range = maxVal - minVal;
  if (range < 0.001) range = 0.001;

  display.setTextSize(1);
  display.setCursor(0, 0);
  display.print("V:");
  display.print(dataPoints[(currentX - 1 + GRAPH_WIDTH) % GRAPH_WIDTH], 2);
  display.setCursor(65, 0);
  display.print("R:");
  display.print(range, 2);

  int baselineY = GRAPH_Y_OFFSET + GRAPH_HEIGHT / 2;
  for (int x = 0; x < GRAPH_WIDTH; x += 10) {
    display.drawPixel(x, baselineY, SSD1306_WHITE);
  }

  for (int i = 0; i < GRAPH_WIDTH - 1; i++) {
    int x1 = i;
    int x2 = i + 1;

    int y1 = GRAPH_Y_OFFSET + GRAPH_HEIGHT - ((dataPoints[i] - minVal) / range * GRAPH_HEIGHT);
    int y2 = GRAPH_Y_OFFSET + GRAPH_HEIGHT - ((dataPoints[x2 % GRAPH_WIDTH] - minVal) / range * GRAPH_HEIGHT);

    y1 = constrain(y1, GRAPH_Y_OFFSET, GRAPH_Y_OFFSET + GRAPH_HEIGHT);
    y2 = constrain(y2, GRAPH_Y_OFFSET, GRAPH_Y_OFFSET + GRAPH_HEIGHT);

    display.drawLine(x1, y1, x2, y2, SSD1306_WHITE);
  }

  int markerY = GRAPH_Y_OFFSET + GRAPH_HEIGHT - ((dataPoints[currentX] - minVal) / range * GRAPH_HEIGHT);
  markerY = constrain(markerY, GRAPH_Y_OFFSET, GRAPH_Y_OFFSET + GRAPH_HEIGHT);
  display.fillCircle(currentX, markerY, 1, SSD1306_WHITE);

  display.display();
}