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

// MAKE CHANGES BELOW:

// Rotary Encoder Inputs:
#define CLK 7
#define DT 8
#define SW 9

#define SCREEN_WIDTH 128 // OLED display width, in pixels
#define SCREEN_HEIGHT 32 // OLED display height, in pixels

// MAKE CHANGES ABOVE:

int currentStateCLK, lastStateCLK;
unsigned long lastButtonPress = 0;

String MenuItem1 = "Speed";
String MenuItem2 = "Axis";

int speedOption;
int SelectedMenuItem;

unsigned long displayTimer;
bool displayTimeEnabled;

bool inMenu = true;
bool inSpeedMenu, inAxisMenu, inJogMenu;

// Declaration for an SSD1306 display connected to I2C (SDA, SCL pins)
#define OLED_RESET     4 // Reset pin # (or -1 if sharing Arduino reset pin)
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

void setup() {
  pinMode(CLK, INPUT);
  pinMode(DT, INPUT);
  pinMode(SW, INPUT_PULLUP);

  Serial.begin(9600);

  // SSD1306_SWITCHCAPVCC = generate display voltage from 3.3V internally
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) { // Address 0x3D for 128x64
    Serial.println(F("SSD1306 allocation failed"));
    for (;;); // Don't proceed, loop forever
  }

  selectMenu(0);

  lastStateCLK = digitalRead(CLK);
}

void loop() {
  while (inMenu) {
    int rotary = checkEncoder();

    if (rotary == 1) {
      SelectedMenuItem++;
      if (SelectedMenuItem > 2) SelectedMenuItem = 0;
      selectMenu(SelectedMenuItem);
    }
    else if (rotary == -1) {
      SelectedMenuItem--;
      if (SelectedMenuItem < 0) SelectedMenuItem = 2;
      selectMenu(SelectedMenuItem);
    }
    else if (rotary == 10) {
      if (SelectedMenuItem == 0) {
        inSpeedMenu = true;
        speedMenu(0);
      }
      else if (SelectedMenuItem == 1) {
        inAxisMenu = true;
        axisMenu(0);
      }

      else if (SelectedMenuItem == 2) {
        inJogMenu = true;
        jogMenu(0);
      }

      inMenu = false;
      SelectedMenuItem = 0;
      delay(300);
    }
    delay(1);
  }

  while (inSpeedMenu) {
    int rotary = checkEncoder();

    if (rotary == 1) {
      SelectedMenuItem++;
      if (SelectedMenuItem > 3) SelectedMenuItem = 0;
      speedMenu(SelectedMenuItem);
    }
    else if (rotary == -1) {
      SelectedMenuItem--;
      if (SelectedMenuItem < 0) SelectedMenuItem = 3;
      speedMenu(SelectedMenuItem);
    }
    else if (rotary == 10) {
      if (SelectedMenuItem == 0) MenuItem1 = "0.025mm";
      else if (SelectedMenuItem == 1) MenuItem1 = "0.25mm";
      else if (SelectedMenuItem == 2) MenuItem1 = "1mm";
      else if (SelectedMenuItem == 3) MenuItem1 = "Fast";
      inSpeedMenu = false;
      inMenu = true;
      SelectedMenuItem = 0;
      selectMenu(SelectedMenuItem);
      delay(300);
    }
    delay(1);
  }

  while (inAxisMenu) {
    int rotary = checkEncoder();

    if (rotary == 1) {
      SelectedMenuItem++;
      if (SelectedMenuItem > 2) SelectedMenuItem = 0;
      axisMenu(SelectedMenuItem);
    }
    else if (rotary == -1) {
      SelectedMenuItem--;
      if (SelectedMenuItem < 0) SelectedMenuItem = 2;
      axisMenu(SelectedMenuItem);
    }
    else if (rotary == 10) {
      if (SelectedMenuItem == 0) MenuItem2 = "X";
      else if (SelectedMenuItem == 1) MenuItem2 = "Y";
      else if (SelectedMenuItem == 2) MenuItem2 = "Z";
      inAxisMenu = false;
      inMenu = true;
      SelectedMenuItem = 0;
      selectMenu(SelectedMenuItem);
      delay(300);
    }
    delay(1);
  }

  while (inJogMenu) {
    int rotary = checkEncoder();

    if (rotary == 1) {
      jogMenu(2); // CW rotation
      displayTimeEnabled = true;
      displayTimer = millis();
    }
    else if (rotary == -1) {
      jogMenu(1); // CCW rotation
      displayTimeEnabled = true;
      displayTimer = millis();
    }
    else if (rotary == 10) {
      inJogMenu = false;
      inMenu = true;
      SelectedMenuItem = 0;
      selectMenu(SelectedMenuItem);
      delay(300);
    }

    if (displayTimeEnabled && displayTimer + 300 < millis()) {
      displayTimeEnabled = false;
      jogMenu(0);
    }
    delay(1);
  }
}

int checkEncoder() {
  int encoder = 0;
  currentStateCLK = digitalRead(CLK);

  if (currentStateCLK != lastStateCLK  && currentStateCLK == 1) {

    if (digitalRead(DT) != currentStateCLK) {   //--
      encoder = 1;
    } else {    // ++
      encoder = -1;
    }
  }

  lastStateCLK = currentStateCLK;

  if (!digitalRead(SW)) {
    if (millis() - lastButtonPress > 50) {
      encoder = 10;
    }

    lastButtonPress = millis();
  }

  return encoder;
}

void selectMenu(int cursor)
{
  display.clearDisplay();

  display.setTextColor(SSD1306_WHITE);
  display.setTextSize(1);

  if (cursor == 0) {
    display.fillRect(0, 0, 42, 32, SSD1306_WHITE);
    display.setTextColor(SSD1306_BLACK);
    if(MenuItem1 == "0.025mm") display.setCursor(0, 12);
    else if(MenuItem1 == "0.25mm") display.setCursor(4, 12);
    else if(MenuItem1 == "1mm") display.setCursor(12, 12);
    else if(MenuItem1 == "Fast") display.setCursor(10, 12);
    else display.setCursor(6, 12);  // Speed
    
    display.print(MenuItem1);

    display.setTextColor(SSD1306_WHITE);
    display.setCursor(52, 12);
    display.print(MenuItem2);
    display.setCursor(98, 12);
    display.print("Jog");
  }

  if (cursor == 1) {
    display.fillRect(42, 0, 43, 32, SSD1306_WHITE);
    display.setTextColor(SSD1306_BLACK);
    display.setCursor(52, 12);
    display.print(MenuItem2);

    display.setTextColor(SSD1306_WHITE);
    if(MenuItem1 == "0.025mm") display.setCursor(0, 12);
    else if(MenuItem1 == "0.25mm") display.setCursor(4, 12);
    else if(MenuItem1 == "1mm") display.setCursor(12, 12);
    else if(MenuItem1 == "Fast") display.setCursor(10, 12);
    else display.setCursor(6, 12);  // Speed
    display.print(MenuItem1);
    display.setCursor(98, 12);
    display.print("Jog");
  }

  if (cursor == 2) {
    display.fillRect(85, 0, 43, 32, SSD1306_WHITE);
    display.setTextColor(SSD1306_BLACK);
    display.setCursor(98, 12);
    display.print("Jog");

    display.setTextColor(SSD1306_WHITE);
    if(MenuItem1 == "0.025mm") display.setCursor(0, 12);
    else if(MenuItem1 == "0.25mm") display.setCursor(4, 12);
    else if(MenuItem1 == "1mm") display.setCursor(12, 12);
    else if(MenuItem1 == "Fast") display.setCursor(10, 12);
    else display.setCursor(6, 12);  // Speed
    display.print(MenuItem1);
    display.setCursor(52, 12);
    display.print(MenuItem2);
  }

  display.display();
}

void speedMenu(int speedoption) {
  display.clearDisplay();

  display.setTextColor(SSD1306_WHITE);
  display.setTextSize(1);

  if (speedoption == 0) {
    display.setCursor(43, 12);
    display.print("0.025mm");
  }
  else if (speedoption == 1) {
    display.setCursor(47, 12);
    display.print("0.25mm");
  }
  else if (speedoption == 2) {
    display.setCursor(56, 12);
    display.print("1mm");
  }
  else if (speedoption == 3) {
    display.setCursor(53, 12);
    display.print("Fast");
  }

  display.display();
}

void axisMenu(int axisoption) {
  display.clearDisplay();

  display.setTextColor(SSD1306_WHITE);
  display.setTextSize(1);
  display.setCursor(62, 12);

  if (axisoption == 0) {
    display.print("X");
  }
  else if (axisoption == 1) {
    display.print("Y");
  }
  else if (axisoption == 2) {
    display.print("Z");
  }

  display.display();
}

void jogMenu(int option) {
  display.clearDisplay();

  display.setTextColor(SSD1306_WHITE);
  display.setTextSize(1);

  if (option == 1) {
    display.setCursor(49, 12);
    display.print("<<<<");
  }
  else if (option == 2) {
    display.setCursor(49, 12);
    display.print(">>>>");
  }
  else if (option == 0) {
    display.setCursor(62, 12);
    display.print("0");
  }

  display.display();
}