#include <U8g2lib.h>
#include <math.h>
#include <FreqMeasure.h>

// Initialize display
U8G2_SSD1306_128X64_NONAME_F_HW_I2C u8g2(U8G2_R0, -1, A5, A4);

float capVoltage = 0, setVoltage = 0;
float previousCapVoltage = -1, previousSetVoltage = -1;
float frequency = 0;

const uint8_t potentiometer = A0, receiverIR = A1, hallEffect = A2;
const uint8_t buttonPin = 7, bleed = 9, dischargePin = 10, chargePin = 11, solenoidPin = 12, indicatorLED = 13;
const uint8_t debounce = 20, holdTime = 1000;

bool mode = false;

int b = 0; // value read from button
int blast = 0; // buffered value of the button's previous state

#define centreX(text) ((128 - u8g2.getStrWidth(text)) / 2)

void setup() {
  Serial.begin(9600);  // Start serial communication
  FreqMeasure.begin();

  pinMode(bleed, OUTPUT);
  pinMode(buttonPin, INPUT);
  pinMode(chargePin, OUTPUT);
  pinMode(dischargePin, OUTPUT);
  pinMode(indicatorLED, OUTPUT);
  pinMode(solenoidPin, OUTPUT);
  
  digitalWrite(buttonPin, HIGH); // pull-up 20k
  digitalWrite(indicatorLED, LOW);
  
  u8g2.begin();  // Initialize the display
  u8g2.clearBuffer();  // Clear the display buffer
  u8g2.setFont(u8g2_font_profont12_tr);  // Set font size
  u8g2.setCursor(30, 20);  // Set cursor position
  u8g2.print(F("James Bray's"));  // Print using F() macro
  u8g2.setFont(u8g2_font_profont17_tr);  // Set font size
  u8g2.setCursor(1, 40);  // Set cursor position
  u8g2.print(F("Marble Machine"));  // Print title
  u8g2.setCursor(34, 60);  // Set cursor position
  u8g2.setFont(u8g2_font_profont10_tr);  // Set smaller font
  u8g2.print(F("2024 11A_EST1"));  // Print using F() macro
  u8g2.sendBuffer();  // Send data to the display
  delay(3000);  // Wait for 3 seconds
  u8g2.clearBuffer();  // Clear buffer
}

void updateOLED() {
  u8g2.clearBuffer();
  u8g2.drawLine(0, 12, 128, 12);
  u8g2.setCursor(0, 10);
  u8g2.setFont(u8g2_font_profont10_tr);

  if (mode) {
    u8g2.print(F("automatic Mode"));
    if (setVoltage > capVoltage) {
      charging();
    } else {
      automatic();
    }
  } else {
    u8g2.print(F("Manual Mode"));
    if (setVoltage > capVoltage) {
      charging();
    } else {
      manual();
    }
  }
  u8g2.sendBuffer();
}

void charging() {
  digitalWrite(chargePin, HIGH);
  u8g2.setFont(u8g2_font_profont12_tr);  // Set font
  u8g2.setCursor(centreX("Set Voltage: 000 V"), 25);
  u8g2.print(F("Set Voltage: "));
  u8g2.print((int)setVoltage);  // Print the integer part only
  u8g2.print(F(" V"));

  u8g2.setCursor(centreX("Cap Voltage: 000 V"), 40);
  u8g2.print(F("Cap Voltage: "));
  
  if (capVoltage < 20) {
    u8g2.print(F("<20 V"));
  } else {
    u8g2.print((int)capVoltage);  // Print the integer part only
    u8g2.print(F(" V"));
  }

  u8g2.setBitmapMode(1);
  u8g2.drawFrame(12, 44, 104, 18);  // Draw voltage frame
  u8g2.drawBox(14, 46, (capVoltage * 100 / setVoltage), 14);  // Draw charging progress bar
}


void automatic(void) {
  digitalWrite(chargePin, LOW);
  discharge();
}

void manual() {
  digitalWrite(chargePin, LOW);
  u8g2.clearBuffer();
  drawHeader();
  u8g2.setFont(u8g2_font_profont12_tr);
  u8g2.setCursor(centreX("Press the button to"), 26);
  u8g2.print(F("Press the button to"));
  u8g2.setCursor(centreX("manually discharge."), 38);
  u8g2.print(F("manually discharge."));
  u8g2.sendBuffer();
  while (digitalRead(buttonPin) == LOW);
  discharge();
}

void discharge() {
  for (int8_t i = 3; i >= 0; i--) {
    u8g2.clearBuffer();
    drawHeader();
    u8g2.setCursor(20, 25);
    u8g2.setFont(u8g2_font_profont12_tr);
    u8g2.print(F("Discharging in:"));
    u8g2.setCursor(56, 50);
    u8g2.setFont(u8g2_font_profont29_tr);
    u8g2.print(i);
    u8g2.sendBuffer();
    delay(1000);  // Wait 1 second
  }
  digitalWrite(solenoidPin, HIGH);
  while (analogRead(receiverIR) < 512);  // Wait for IR signal
  digitalWrite(solenoidPin, LOW);
  digitalWrite(dischargePin, HIGH);
  delay(1000);
  digitalWrite(dischargePin, LOW);
}

void drawHeader() {
  u8g2.setFont(u8g2_font_profont10_tr);
  u8g2.drawLine(0, 12, 128, 12);
  u8g2.setCursor(0, 10);
  if (mode) {
    u8g2.print(F("automatic Mode"));
  } else {
    u8g2.print(F("Manual Mode"));
  }
  u8g2.setCursor(0, 100);
  u8g2.print(capVoltage);
}

void checkButton() {
  long bdown; // time the button was pressed down
  long bup; // time the button was released
  boolean ignoreUp = false; // whether to ignore the button release because the click+hold was triggered

  b = digitalRead(buttonPin); // read the state of the button
  if (b == LOW && blast == HIGH && (millis() - bup) > long(debounce)) { // test for button pressed and store the down time
    bdown = millis();
  }
  
  if (b == HIGH && blast == LOW && (millis() - bdown) > long(debounce)) { // test for button release and store the up time
    if (ignoreUp == false) {
      mode = !mode;
    } else {
      ignoreUp = false;
    }
    bup = millis();
  }
  
  if (b == LOW && (millis() - bdown) > long(holdTime)) { // test for button held down for longer than the hold time
    manual();
    ignoreUp = true;
    bdown = millis();
  }
  blast = b;
}

void loop() {
  delay(1);

  // Local variables within loop
  int16_t potValue = analogRead(potentiometer);  // Read the potentiometer value
  static double sum = 0;
  static int count = 0;

  if (FreqMeasure.available()) {
    sum += FreqMeasure.read();
    count++;
    if (count > 10) {
      frequency = FreqMeasure.countToFrequency(sum / count);
      sum = 0;
      count = 0;
    }
  }

  setVoltage = map(potValue, 10, 1013, 20, 80) * 5;  // Map potentiometer value to voltage
  capVoltage = (47.05 * exp(0.002425*frequency)) -37.87;  // Map capacitor value to voltage (non-linear, calculated experimentally with line-of-best-fit)

  checkButton();

  // Only update display if there is a change in values or mode
  if (capVoltage != previousCapVoltage || setVoltage != previousSetVoltage || b != blast) {
    updateOLED();
    previousCapVoltage = capVoltage;
    previousSetVoltage = setVoltage;
  }
}