#include <Wire.h>
#include <OneWire.h>
#include <DallasTemperature.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>

// Sensors
#define ONE_WIRE_BUS 4
#define GSR_PIN 34
#define HEART_PIN 35

// Output LEDs
#define COOLING_LED 18
#define VIBRATION_LED 21

// OLED connections
#define OLED_SDA 25
#define OLED_SCL 26
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

OneWire oneWire(ONE_WIRE_BUS);
DallasTemperature sensors(&oneWire);

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

const int RR_COUNT = 6;

// Simulated intervals between heartbeats (ms)
float rrIntervals[RR_COUNT] = {
  800, 820, 790, 810, 830, 805
};

float calculateHRV() {
  float sum = 0;

  for (int i = 0; i < RR_COUNT - 1; i++) {
    float difference =
      rrIntervals[i + 1] - rrIntervals[i];

    sum += difference * difference;
  }

  float mean = sum / (RR_COUNT - 1);

  return sqrt(mean);
}

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

  sensors.begin();

  pinMode(COOLING_LED, OUTPUT);
  pinMode(VIBRATION_LED, OUTPUT);

  digitalWrite(COOLING_LED, LOW);
  digitalWrite(VIBRATION_LED, LOW);

  Wire.begin(OLED_SDA, OLED_SCL);

  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println("OLED Error!");
  } else {
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 0);
    display.println("Migraine Headband");
    display.println("System Starting...");
    display.display();
  }

  delay(1500);
}

void loop() {
  int gsrValue = analogRead(GSR_PIN);
float hrvValue = calculateHRV();
  int heartSignal = analogRead(HEART_PIN);
  int heartRate = map(heartSignal, 0, 4095, 60, 120);

  sensors.requestTemperatures();
  float temperature = sensors.getTempCByIndex(0);

  bool tempValid =
    temperature != DEVICE_DISCONNECTED_C;

  bool highStatus =
    gsrValue >= 2800 && heartRate >= 100;

  // Reset outputs each cycle
  digitalWrite(COOLING_LED, LOW);
  digitalWrite(VIBRATION_LED, LOW);

  if (tempValid) {
    if (temperature >= 35.0) {
      digitalWrite(COOLING_LED, HIGH);
    }
    else if (temperature >= 30.0 || highStatus) {
      digitalWrite(VIBRATION_LED, HIGH);
    }
  }

  // Serial Monitor
  Serial.print("GSR: ");
  Serial.println(gsrValue);

  Serial.print("Heart Rate: ");
  Serial.println(heartRate);

  Serial.print("Temperature: ");
  Serial.println(temperature);

  Serial.print("Status: ");
  if (!tempValid) {
    Serial.println("SENSOR ERROR");
  } else {
    Serial.println(highStatus ? "HIGH" : "NORMAL");
  }

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

  display.setCursor(0, 0);
  display.println("MIGRAINE HEADBAND");

  display.setCursor(0, 16);
  display.print("HR: ");
  display.print(heartRate);
  display.println(" BPM (SIM)");

  display.setCursor(0, 26);
  display.print("GSR: ");
  display.println(gsrValue);

  display.setCursor(0, 36);
  display.print("Temp: ");
  if (tempValid) {
    display.print(temperature, 1);
    display.println(" C");
  } else {
    display.println("ERROR");
  }

  display.setCursor(0, 48);
  display.print("Status: ");
  if (!tempValid) {
    display.println("ERROR");
  } else {
    display.println(highStatus ? "HIGH" : "NORMAL");
  }
display.setCursor(0, 56);
display.print("HRV: ");
display.print(hrvValue, 1);
display.print(" ms SIM");
  display.display();

  delay(1000);
}