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

#ifndef BIOMED_DIAGNOSTIC
#define BIOMED_DIAGNOSTIC false
#endif

// Reviewed learning-lab pin plan for the ESP32 DevKit-style board in this kit.
constexpr uint8_t PIN_ECG_OUT = 34;
constexpr uint8_t PIN_ECG_LO_PLUS = 32;
constexpr uint8_t PIN_ECG_LO_MINUS = 33;
constexpr uint8_t PIN_I2C_SDA = 21;
constexpr uint8_t PIN_I2C_SCL = 22;

constexpr uint32_t SERIAL_BAUD = 115200;
constexpr uint32_t ECG_PERIOD_US = 4000;  // 250 samples/second.
constexpr uint32_t DISPLAY_PERIOD_MS = 200;
constexpr uint32_t STATUS_PERIOD_MS = 1000;
constexpr uint32_t FINGER_IR_MIN = 50000;  // Provisional; tune with the exact module.
constexpr uint8_t OLED_WIDTH = 128;
constexpr uint8_t OLED_HEIGHT = 64;
constexpr int8_t OLED_RESET = -1;

enum class RunMode : uint8_t {
  Paused,
  Optical,
  Ecg,
  Combined,
  PlotEcg
};

Adafruit_SSD1306 display(OLED_WIDTH, OLED_HEIGHT, &Wire, OLED_RESET);
MAX30105 opticalSensor;

RunMode runMode = RunMode::Paused;
bool safetyAcknowledged = false;
bool oledReady = false;
bool opticalReady = false;
uint8_t oledAddress = 0;

uint32_t nextEcgUs = 0;
uint32_t nextDisplayMs = 0;
uint32_t nextStatusMs = 0;
uint32_t lastBeatMs = 0;
uint32_t lastOpticalSampleMs = 0;
uint32_t latestRed = 0;
uint32_t latestIr = 0;
uint16_t latestEcg = 0;
bool latestLeadOffPlus = true;
bool latestLeadOffMinus = true;
float latestBpm = NAN;

float bpmHistory[4] = {NAN, NAN, NAN, NAN};
uint8_t bpmIndex = 0;
uint8_t bpmCount = 0;
uint8_t combinedEcgDivider = 0;
uint8_t opticalPrintDivider = 0;

String commandBuffer;

const char *modeName(RunMode mode) {
  switch (mode) {
    case RunMode::Optical: return "OPTICAL";
    case RunMode::Ecg: return "ECG";
    case RunMode::Combined: return "COMBINED";
    case RunMode::PlotEcg: return "PLOT_ECG";
    default: return "PAUSED";
  }
}

bool isEcgMode(RunMode mode) {
  return mode == RunMode::Ecg || mode == RunMode::Combined || mode == RunMode::PlotEcg;
}

bool isOpticalMode(RunMode mode) {
  return mode == RunMode::Optical || mode == RunMode::Combined;
}

bool i2cResponds(uint8_t address) {
  Wire.beginTransmission(address);
  return Wire.endTransmission() == 0;
}

void printSafety() {
  Serial.println();
  Serial.println(F("SAFETY - EDUCATION ONLY, NOT A MEDICAL DEVICE"));
  Serial.println(F("Never use this kit for diagnosis, treatment, emergencies, patient monitoring, or health decisions."));
  Serial.println(F("Before body-connected ECG use: run the laptop from battery only and disconnect its charger, dock,"));
  Serial.println(F("mains-connected peripherals, and test instruments. USB is not medical isolation."));
  Serial.println(F("Keep electrodes off the body while wiring, uploading, resetting, or changing connections."));
  Serial.println(F("Use only on clean, intact skin; stop if irritation, discomfort, or unexpected heating occurs."));
  Serial.println(F("Type I-UNDERSTAND to unlock ECG modes for this power session."));
  Serial.println();
}

void printHelp() {
  Serial.println(F("Commands:"));
  Serial.println(F("  help          - show this list"));
  Serial.println(F("  safety        - repeat the safety notice"));
  Serial.println(F("  scan          - scan the I2C bus"));
  Serial.println(F("  status        - report sensors, pins, and current mode"));
  Serial.println(F("  optical       - MAX30102 raw red/IR plus experimental pulse interval"));
  Serial.println(F("  ecg           - structured AD8232 raw samples; safety acknowledgement required"));
  Serial.println(F("  plot-ecg      - Arduino Serial Plotter output; safety acknowledgement required"));
  Serial.println(F("  combined      - reduced-rate ECG and optical stream; acknowledgement required"));
  Serial.println(F("  pause         - stop acquisition and return to the safety screen"));
  Serial.println(F("No SpO2 value is calculated. BPM is experimental and is shown only with adequate optical signal."));
}

void scanI2c() {
  uint8_t count = 0;
  Serial.println(F("I2C_SCAN_BEGIN"));
  for (uint8_t address = 1; address < 127; ++address) {
    if (i2cResponds(address)) {
      Serial.print(F("I2C_DEVICE,0x"));
      if (address < 16) Serial.print('0');
      Serial.println(address, HEX);
      ++count;
    }
  }
  Serial.print(F("I2C_SCAN_END,count="));
  Serial.println(count);
}

void showLines(const __FlashStringHelper *line1,
               const __FlashStringHelper *line2,
               const __FlashStringHelper *line3 = nullptr,
               const __FlashStringHelper *line4 = nullptr) {
  if (!oledReady) return;
  display.clearDisplay();
  display.setTextColor(SSD1306_WHITE);
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.println(line1);
  display.println(line2);
  if (line3 != nullptr) display.println(line3);
  if (line4 != nullptr) display.println(line4);
  display.display();
}

void initDisplay() {
  if (BIOMED_DIAGNOSTIC) {
    oledReady = false;
    return;
  }

  if (i2cResponds(0x3C)) oledAddress = 0x3C;
  else if (i2cResponds(0x3D)) oledAddress = 0x3D;
  else return;

  oledReady = display.begin(SSD1306_SWITCHCAPVCC, oledAddress);
  if (oledReady) {
    showLines(F("Biomedical Lab"), F("EDUCATION ONLY"), F("Type help in Serial"));
  }
}

void initOpticalSensor() {
  if (BIOMED_DIAGNOSTIC) {
    opticalReady = true;
    return;
  }

  opticalReady = opticalSensor.begin(Wire, I2C_SPEED_FAST);
  if (!opticalReady) return;

  // Red + IR at 100 samples/s. Values are for learning and must be tuned on the exact breakout.
  opticalSensor.setup(0x1F, 4, 2, 100, 411, 4096);
  opticalSensor.setPulseAmplitudeGreen(0);
  opticalSensor.clearFIFO();
}

void printStatus() {
  Serial.print(F("STATUS,mode="));
  Serial.print(modeName(runMode));
  Serial.print(F(",diagnostic="));
  Serial.print(BIOMED_DIAGNOSTIC ? 1 : 0);
  Serial.print(F(",safety_ack="));
  Serial.print(safetyAcknowledged ? 1 : 0);
  Serial.print(F(",oled="));
  Serial.print(oledReady ? 1 : 0);
  Serial.print(F(",oled_addr="));
  if (oledAddress == 0) Serial.print(F("NA"));
  else {
    Serial.print(F("0x"));
    Serial.print(oledAddress, HEX);
  }
  Serial.print(F(",max30102="));
  Serial.print(opticalReady ? 1 : 0);
  Serial.print(F(",ecg_pin="));
  Serial.print(PIN_ECG_OUT);
  Serial.print(F(",lo_plus_pin="));
  Serial.print(PIN_ECG_LO_PLUS);
  Serial.print(F(",lo_minus_pin="));
  Serial.println(PIN_ECG_LO_MINUS);
}

void resetBpm() {
  latestBpm = NAN;
  bpmIndex = 0;
  bpmCount = 0;
  lastBeatMs = 0;
  for (float &entry : bpmHistory) entry = NAN;
}

void setMode(RunMode requested) {
  if (isEcgMode(requested) && !safetyAcknowledged) {
    Serial.println(F("DENIED: Read the safety notice and type I-UNDERSTAND first."));
    return;
  }
  if (isOpticalMode(requested) && !opticalReady) {
    Serial.println(F("DENIED: MAX30102 was not detected at I2C address 0x57."));
    return;
  }

  runMode = requested;
  nextEcgUs = micros();
  nextDisplayMs = millis();
  combinedEcgDivider = 0;
  opticalPrintDivider = 0;
  if (!isOpticalMode(runMode)) resetBpm();

  Serial.print(F("MODE,"));
  Serial.println(modeName(runMode));
  if (runMode == RunMode::Optical) {
    Serial.println(F("PPG_HEADER,t_ms,red,ir,bpm_estimate,signal"));
  } else if (runMode == RunMode::Ecg) {
    Serial.println(F("ECG_HEADER,t_us,raw,lead_off_plus,lead_off_minus"));
  } else if (runMode == RunMode::Combined) {
    Serial.println(F("COMBINED: ECG is reduced to 50 Hz and PPG to about 25 Hz for serial bandwidth."));
  } else if (runMode == RunMode::PlotEcg) {
    Serial.println(F("Serial Plotter channels: ECG_RAW and LEADS_OFF. Disconnect all mains-connected equipment."));
  }
}

void handleCommand(String command) {
  command.trim();
  command.toLowerCase();
  if (command.length() == 0) return;

  if (command == "i-understand") {
    safetyAcknowledged = true;
    Serial.println(F("SAFETY_ACKNOWLEDGED for this power session. This does not verify the wiring or make the kit medical equipment."));
  } else if (command == "help") {
    printHelp();
  } else if (command == "safety") {
    printSafety();
  } else if (command == "scan") {
    scanI2c();
  } else if (command == "status") {
    printStatus();
  } else if (command == "optical") {
    setMode(RunMode::Optical);
  } else if (command == "ecg") {
    setMode(RunMode::Ecg);
  } else if (command == "combined") {
    setMode(RunMode::Combined);
  } else if (command == "plot-ecg") {
    setMode(RunMode::PlotEcg);
  } else if (command == "pause") {
    setMode(RunMode::Paused);
  } else {
    Serial.print(F("UNKNOWN_COMMAND,"));
    Serial.println(command);
    printHelp();
  }
}

void serviceSerial() {
  while (Serial.available() > 0) {
    const char incoming = static_cast<char>(Serial.read());
    if (incoming == '\n' || incoming == '\r') {
      if (commandBuffer.length() > 0) {
        handleCommand(commandBuffer);
        commandBuffer = "";
      }
    } else if (commandBuffer.length() < 48) {
      commandBuffer += incoming;
    }
  }
}

void readLeadState() {
  if (BIOMED_DIAGNOSTIC) {
    latestLeadOffPlus = false;
    latestLeadOffMinus = false;
    return;
  }
  latestLeadOffPlus = digitalRead(PIN_ECG_LO_PLUS) == HIGH;
  latestLeadOffMinus = digitalRead(PIN_ECG_LO_MINUS) == HIGH;
}

uint16_t diagnosticEcgSample() {
  const uint16_t phase = (millis() / 4) % 250;
  if (phase < 8) return 2100 + phase * 180;
  if (phase < 18) return 3540 - (phase - 8) * 260;
  if (phase < 28) return 940 + (phase - 18) * 120;
  return 2048 + ((phase % 20) - 10) * 3;
}

void serviceEcg() {
  if (!isEcgMode(runMode)) return;

  const uint32_t nowUs = micros();
  if (static_cast<int32_t>(nowUs - nextEcgUs) < 0) return;
  nextEcgUs += ECG_PERIOD_US;
  if (static_cast<int32_t>(nowUs - nextEcgUs) > static_cast<int32_t>(ECG_PERIOD_US * 4)) {
    nextEcgUs = nowUs + ECG_PERIOD_US;
  }

  readLeadState();
  const bool leadsOff = latestLeadOffPlus || latestLeadOffMinus;
  latestEcg = BIOMED_DIAGNOSTIC ? diagnosticEcgSample() : analogRead(PIN_ECG_OUT);

  if (runMode == RunMode::PlotEcg) {
    Serial.print(F("ECG_RAW:"));
    Serial.print(leadsOff ? 0 : latestEcg);
    Serial.print(F(",LEADS_OFF:"));
    Serial.println(leadsOff ? 4095 : 0);
    return;
  }

  if (runMode == RunMode::Combined && ++combinedEcgDivider < 5) return;
  combinedEcgDivider = 0;

  Serial.print(F("ECG,"));
  Serial.print(nowUs);
  Serial.print(',');
  if (leadsOff) Serial.print(F("NA"));
  else Serial.print(latestEcg);
  Serial.print(',');
  Serial.print(latestLeadOffPlus ? 1 : 0);
  Serial.print(',');
  Serial.println(latestLeadOffMinus ? 1 : 0);
}

void addBpm(float bpm) {
  bpmHistory[bpmIndex] = bpm;
  bpmIndex = (bpmIndex + 1) % 4;
  if (bpmCount < 4) ++bpmCount;

  float sum = 0.0f;
  for (uint8_t i = 0; i < bpmCount; ++i) sum += bpmHistory[i];
  latestBpm = sum / bpmCount;
}

void processOpticalSample(uint32_t red, uint32_t ir, uint32_t nowMs) {
  latestRed = red;
  latestIr = ir;
  lastOpticalSampleMs = nowMs;

  const bool signalPresent = ir >= FINGER_IR_MIN;
  if (!signalPresent) {
    resetBpm();
  } else if (checkForBeat(static_cast<int32_t>(ir))) {
    if (lastBeatMs != 0) {
      const uint32_t interval = nowMs - lastBeatMs;
      if (interval > 0) {
        const float bpm = 60000.0f / interval;
        if (bpm >= 30.0f && bpm <= 220.0f) addBpm(bpm);
      }
    }
    lastBeatMs = nowMs;
  }

  if (++opticalPrintDivider < 4) return;
  opticalPrintDivider = 0;
  Serial.print(F("PPG,"));
  Serial.print(nowMs);
  Serial.print(',');
  Serial.print(red);
  Serial.print(',');
  Serial.print(ir);
  Serial.print(',');
  if (isnan(latestBpm)) Serial.print(F("NA"));
  else Serial.print(latestBpm, 1);
  Serial.print(',');
  Serial.println(signalPresent ? F("PRESENT") : F("LOW"));
}

void serviceOptical() {
  if (!isOpticalMode(runMode) || !opticalReady) return;

  if (BIOMED_DIAGNOSTIC) {
    static uint32_t nextSampleMs = 0;
    const uint32_t nowMs = millis();
    if (static_cast<int32_t>(nowMs - nextSampleMs) < 0) return;
    nextSampleMs = nowMs + 10;
    const uint32_t phase = nowMs % 800;
    const uint32_t pulse = phase < 80 ? (80 - phase) * 700 : 0;
    processOpticalSample(62000 + pulse / 2, 90000 + pulse, nowMs);
    return;
  }

  opticalSensor.check();
  while (opticalSensor.available()) {
    const uint32_t red = opticalSensor.getFIFORed();
    const uint32_t ir = opticalSensor.getFIFOIR();
    opticalSensor.nextSample();
    processOpticalSample(red, ir, millis());
  }
}

void updateDisplay() {
  if (!oledReady) return;
  const uint32_t nowMs = millis();
  if (static_cast<int32_t>(nowMs - nextDisplayMs) < 0) return;
  nextDisplayMs = nowMs + DISPLAY_PERIOD_MS;

  display.clearDisplay();
  display.setTextColor(SSD1306_WHITE);
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.print(F("MODE: "));
  display.println(modeName(runMode));
  display.println(F("EDUCATION ONLY"));

  if (runMode == RunMode::Paused) {
    display.println(F("Type help in Serial"));
    display.println(F("ECG requires safety ack"));
  } else if (isEcgMode(runMode)) {
    const bool leadsOff = latestLeadOffPlus || latestLeadOffMinus;
    display.print(F("ECG: "));
    if (leadsOff) display.println(F("LEADS OFF"));
    else display.println(latestEcg);
    display.print(F("LO+ "));
    display.print(latestLeadOffPlus ? 1 : 0);
    display.print(F("  LO- "));
    display.println(latestLeadOffMinus ? 1 : 0);
  }

  if (isOpticalMode(runMode)) {
    display.print(F("IR: "));
    display.println(latestIr);
    display.print(F("BPM est: "));
    if (isnan(latestBpm)) display.println(F("--"));
    else display.println(latestBpm, 1);
    display.println(F("No SpO2 calculation"));
  }
  display.display();
}

void setup() {
  Serial.begin(SERIAL_BAUD);
  delay(300);
  commandBuffer.reserve(48);

  pinMode(PIN_ECG_LO_PLUS, INPUT);
  pinMode(PIN_ECG_LO_MINUS, INPUT);
  analogReadResolution(12);
  analogSetPinAttenuation(PIN_ECG_OUT, ADC_11db);

  Wire.begin(PIN_I2C_SDA, PIN_I2C_SCL);
  Wire.setClock(400000);
  initDisplay();
  initOpticalSensor();

  Serial.println(F("Biomedical Sensor Learning Lab firmware"));
  Serial.println(F("Compiler-reviewed firmware is not a physical safety or performance test."));
  printSafety();
  printHelp();
  printStatus();
}

void loop() {
  serviceSerial();
  serviceEcg();
  serviceOptical();
  updateDisplay();

  const uint32_t nowMs = millis();
  if (runMode != RunMode::Paused && static_cast<int32_t>(nowMs - nextStatusMs) >= 0) {
    nextStatusMs = nowMs + STATUS_PERIOD_MS;
    if (isOpticalMode(runMode) && !BIOMED_DIAGNOSTIC && nowMs - lastOpticalSampleMs > 1500) {
      Serial.println(F("WARN,MAX30102_NO_NEW_SAMPLES"));
    }
  }
}