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

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define SCREEN_ADDRESS 0x3C

// Buttons are wired between the pin and GND, so a press reads LOW
#define BTN_UP    26  // menu up / vertical line longer
#define BTN_DOWN  2   // menu down / vertical line shorter
#define BTN_LEFT  27  // bottom line shorter
#define BTN_RIGHT 1   // bottom line longer
#define BTN_A     3   // start / submit / next
#define BTN_B     4   // back to menu

// Muller-Lyer layout: two stacked lines centred on the screen
#define ML_CX       (SCREEN_WIDTH / 2)
#define ML_TOP_Y    20   // reference line
#define ML_BOTTOM_Y 44   // adjustable line
#define ML_FIN      6    // fin length along each axis (45 deg)

// Muller-Lyer lengths (kept even so both ends sit on whole pixels)
#define ML_MIN      16
#define ML_MAX      110
#define ML_REF_MIN  40
#define ML_REF_MAX  90
#define ML_STEP     2    // each press moves both ends 1 px

// Vertical-Horizontal layout: inverted T on the left, text on the right
#define VH_CX       36   // x of the vertical line
#define VH_BASE_Y   58   // y of the horizontal line
#define VH_TEXT_X   80

// Vertical-Horizontal lengths
#define VH_MIN      8
#define VH_MAX      54   // vertical top can reach y = 4
#define VH_REF_MIN  20
#define VH_REF_MAX  48
#define VH_STEP     1

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

enum Mode { MULLER_LYER, VERT_HORIZ };
enum State { MENU, ADJUST, RESULT };

Mode mode = MULLER_LYER;
State state = MENU;
int refLen;          // fixed line
int adjLen;          // player-controlled line
bool refInward;      // Muller-Lyer: fin style of the top line; bottom gets the opposite
int rounds[2];       // per mode
float totalAccuracy[2];

// ---------- drawing ----------

void printCentered(const char *text, int y) {
  display.setCursor((SCREEN_WIDTH - strlen(text) * 6) / 2, y);
  display.print(text);
}

void printAt(int x, int y, const char *text) {
  display.setCursor(x, y);
  display.print(text);
}

// Horizontal line of length len centred on the screen. Inward fins are
// arrowheads (<-->) and make the line look shorter; outward fins are
// tails (>--<) and make it look longer.
void drawArrowFigure(int y, int len, bool inward, bool fins) {
  int x1 = ML_CX - len / 2;
  int x2 = ML_CX + len / 2;
  display.drawFastHLine(x1, y, len + 1, SSD1306_WHITE);
  if (!fins) return;

  int d = inward ? ML_FIN : -ML_FIN;
  display.drawLine(x1, y, x1 + d, y - ML_FIN, SSD1306_WHITE);
  display.drawLine(x1, y, x1 + d, y + ML_FIN, SSD1306_WHITE);
  display.drawLine(x2, y, x2 - d, y - ML_FIN, SSD1306_WHITE);
  display.drawLine(x2, y, x2 - d, y + ML_FIN, SSD1306_WHITE);
}

// Inverted T: a vertical line rising from the middle of a horizontal one.
// The vertical tends to look longer than a horizontal of equal length.
void drawTFigure(int horizLen, int vertLen) {
  display.drawFastHLine(VH_CX - horizLen / 2, VH_BASE_Y, horizLen + 1, SSD1306_WHITE);
  display.drawFastVLine(VH_CX, VH_BASE_Y - vertLen, vertLen + 1, SSD1306_WHITE);
}

void drawMenu() {
  display.clearDisplay();
  printCentered("OPTICAL ILLUSIONS", 0);
  printAt(0, 16, mode == MULLER_LYER ? "> Muller-Lyer" : "  Muller-Lyer");
  printAt(0, 28, mode == VERT_HORIZ ? "> Vertical-Horizontal" : "  Vertical-Horizontal");
  printCentered("U/D pick  A start", 46);
  printCentered("B: back to menu", 56);
  display.display();
}

void drawRound() {
  char line[22];
  snprintf(line, sizeof line, "Round %d", rounds[mode] + 1);

  display.clearDisplay();

  if (mode == MULLER_LYER) {
    printCentered(line, 0);
    drawArrowFigure(ML_TOP_Y, refLen, refInward, true);
    drawArrowFigure(ML_BOTTOM_Y, adjLen, !refInward, true);
    printCentered("L/R adjust  A submit", 56);
  } else {
    drawTFigure(refLen, adjLen);
    printAt(VH_TEXT_X, 0, line);
    printAt(VH_TEXT_X, 16, "Match");
    printAt(VH_TEXT_X, 26, "the base");
    printAt(VH_TEXT_X, 42, "U/D");
    printAt(VH_TEXT_X, 52, "A submit");
  }

  display.display();
}

void drawResult(int diff, float accuracy) {
  char line[22];
  int acc = (int)(accuracy + 0.5);
  int avg = (int)(totalAccuracy[mode] / rounds[mode] + 0.5);

  display.clearDisplay();

  if (mode == MULLER_LYER) {
    snprintf(line, sizeof line, "Acc %d%%   Avg %d%%", acc, avg);
    printCentered(line, 0);

    // Show the true lengths with the fins removed, plus dotted guides
    // dropping from the ends of the reference line
    drawArrowFigure(ML_TOP_Y, refLen, refInward, false);
    drawArrowFigure(ML_BOTTOM_Y, adjLen, !refInward, false);
    for (int y = ML_TOP_Y + 3; y < ML_BOTTOM_Y; y += 3) {
      display.drawPixel(ML_CX - refLen / 2, y, SSD1306_WHITE);
      display.drawPixel(ML_CX + refLen / 2, y, SSD1306_WHITE);
    }

    snprintf(line, sizeof line, "Off by %+dpx  A: next", diff);
    printCentered(line, 56);
  } else {
    // Dotted line at the height the vertical should have reached
    drawTFigure(refLen, adjLen);
    int targetY = VH_BASE_Y - refLen;
    for (int x = VH_CX - refLen / 2; x <= VH_CX + refLen / 2; x += 3)
      display.drawPixel(x, targetY, SSD1306_WHITE);

    snprintf(line, sizeof line, "Acc %d%%", acc);
    printAt(VH_TEXT_X, 0, line);
    snprintf(line, sizeof line, "Avg %d%%", avg);
    printAt(VH_TEXT_X, 10, line);
    printAt(VH_TEXT_X, 24, "Off by");
    snprintf(line, sizeof line, "%+dpx", diff);
    printAt(VH_TEXT_X, 34, line);
    printAt(VH_TEXT_X, 52, "A: next");
  }

  display.display();
}

// ---------- input ----------

// True once per press (on the press edge), ignoring contact bounce
bool justPressed(uint8_t pin) {
  static bool wasDown[30];
  static unsigned long lastChange[30];
  bool down = digitalRead(pin) == LOW;
  if (down != wasDown[pin] && millis() - lastChange[pin] > 30) {
    wasDown[pin] = down;
    lastChange[pin] = millis();
    return down;
  }
  return false;
}

// True once on press, then repeatedly while held (like a keyboard key)
bool pressedOrHeld(uint8_t pin) {
  static unsigned long nextFire[30];
  static unsigned long lastDown[30];
  unsigned long now = millis();

  if (digitalRead(pin) == HIGH) {
    if (now - lastDown[pin] > 30) nextFire[pin] = 0;  // released, not bouncing
    return false;
  }

  lastDown[pin] = now;
  if (nextFire[pin] == 0) {          // first press
    nextFire[pin] = now + 350;
    return true;
  }
  if (now >= nextFire[pin]) {        // held: repeat
    nextFire[pin] = now + 40;
    return true;
  }
  return false;
}

// ---------- game logic ----------

void showMenu() {
  state = MENU;
  drawMenu();
}

void newRound() {
  // Start the adjustable line clearly wrong so the player has to work
  if (mode == MULLER_LYER) {
    refLen = random(ML_REF_MIN / 2, ML_REF_MAX / 2 + 1) * 2;
    refInward = random(2);
    do {
      adjLen = random(ML_MIN / 2, ML_MAX / 2 + 1) * 2;
    } while (abs(adjLen - refLen) < 20);
  } else {
    refLen = random(VH_REF_MIN / 2, VH_REF_MAX / 2 + 1) * 2;
    do {
      adjLen = random(VH_MIN, VH_MAX + 1);
    } while (abs(adjLen - refLen) < 10);
  }

  state = ADJUST;
  drawRound();
}

void submit() {
  int diff = adjLen - refLen;
  float accuracy = 100.0 * (1.0 - (float)abs(diff) / refLen);
  if (accuracy < 0) accuracy = 0;

  rounds[mode]++;
  totalAccuracy[mode] += accuracy;

  if (mode == MULLER_LYER) {
    Serial.print("[Muller-Lyer] ref=");
    Serial.print(refLen);
    Serial.print(refInward ? " (arrowheads)" : " (tails)");
  } else {
    Serial.print("[Vertical-Horizontal] horizontal=");
    Serial.print(refLen);
  }
  Serial.print(" guess=");
  Serial.print(adjLen);
  Serial.print(" diff=");
  Serial.println(diff);

  state = RESULT;
  drawResult(diff, accuracy);
}

// ---------- main ----------

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

  pinMode(BTN_UP, INPUT_PULLUP);
  pinMode(BTN_DOWN, INPUT_PULLUP);
  pinMode(BTN_LEFT, INPUT_PULLUP);
  pinMode(BTN_RIGHT, INPUT_PULLUP);
  pinMode(BTN_A, INPUT_PULLUP);
  pinMode(BTN_B, INPUT_PULLUP);

  Wire1.setSDA(6);
  Wire1.setSCL(7);
  Wire1.begin();

  if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    Serial.println("SSD1306 failed");
    while (true);
  }

  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setTextWrap(false);
  showMenu();
}

void loop() {
  switch (state) {
    case MENU:
      if (justPressed(BTN_UP) || justPressed(BTN_DOWN)) {
        mode = (mode == MULLER_LYER) ? VERT_HORIZ : MULLER_LYER;
        drawMenu();
      }
      if (justPressed(BTN_A)) {
        randomSeed(micros());  // exact time of the button press is a good seed
        newRound();
      }
      break;

    case ADJUST: {
      int before = adjLen;
      if (mode == MULLER_LYER) {
        if (pressedOrHeld(BTN_LEFT)  && adjLen > ML_MIN) adjLen -= ML_STEP;
        if (pressedOrHeld(BTN_RIGHT) && adjLen < ML_MAX) adjLen += ML_STEP;
      } else {
        if (pressedOrHeld(BTN_DOWN) && adjLen > VH_MIN) adjLen -= VH_STEP;
        if (pressedOrHeld(BTN_UP)   && adjLen < VH_MAX) adjLen += VH_STEP;
      }
      if (adjLen != before) drawRound();

      if (justPressed(BTN_A)) submit();
      else if (justPressed(BTN_B)) showMenu();
      break;
    }

    case RESULT:
      if (justPressed(BTN_A)) newRound();
      else if (justPressed(BTN_B)) showMenu();
      break;
  }
}