// Reaction Time Game using 4017, OLED (128x32), and Serial Debugging
#include <SPI.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>

// OLED Display config
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 32
#define OLED_RESET    -1
#define SCREEN_ADDRESS 0x3C
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// 4017 control pins
const int CLOCK_PIN = 10;
const int RESET_PIN = 9;

// Input pin (shared button output)
const int BUTTON_INPUT_PIN = 12;

// Touch sensor input pin (used instead of pin 7 start)
const int TOUCH_SENSOR_PIN = 7;

// Constants
const int NUM_ROUNDS = 5;

// Variables
unsigned long reactionTimes[NUM_ROUNDS];
int currentRound = 0;

// LED positions on 4017
const int ledSteps[4] = {0, 1, 2, 3};

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

  // OLED setup
  if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    Serial.println(F("SSD1306 allocation failed"));
    while (true);
  }
  display.clearDisplay();
  display.display();

  // 4017 control pins
  pinMode(CLOCK_PIN, OUTPUT);
  pinMode(RESET_PIN, OUTPUT);
  digitalWrite(CLOCK_PIN, LOW);
  digitalWrite(RESET_PIN, LOW);

  // Inputs
  pinMode(BUTTON_INPUT_PIN, INPUT);
  pinMode(TOUCH_SENSOR_PIN, INPUT);

  Serial.println("Setup complete.");
  showStartScreen();
}

void loop() {
  Serial.println("Waiting for touch sensor to start the game...");
  while (digitalRead(TOUCH_SENSOR_PIN) == LOW) {
    // wait
  }
  Serial.println("Touch detected! Starting game...");
  runCountdown();
  runGame();
  showResults();
  showStartScreen();
  delay(2000);
}

void showStartScreen() {
  display.clearDisplay();
  display.setTextSize(2);
  display.setTextColor(WHITE);
  display.setCursor(0, 0);
  display.println("Touch Pads");
  display.println("to START");
  display.display();
  Serial.println("Displayed start screen, waiting for touch sensor...");
  Serial.println("---------------------");
}

void runCountdown() {
  display.clearDisplay();
  display.setTextSize(2);
  display.setTextColor(WHITE);

  for (int i = 3; i > 0; i--) {
    display.clearDisplay();
    display.setCursor(55, 5);
    display.print(i);
    display.display();
    Serial.print("Countdown: "); Serial.println(i);
    delay(1000);
  }

  display.clearDisplay();
  display.setCursor(40, 5);
  display.print("GO!");
  display.display();
  Serial.println("GO!");
  delay(1000);
  display.clearDisplay();
}

void runGame() {
  Serial.println("Running game for 5 rounds...");
  for (currentRound = 0; currentRound < NUM_ROUNDS; currentRound++) {

    if (currentRound > 0) {  // After first round, wait for touch sensor & countdown
      waitForTouchToContinue();
    }

    Serial.println("---------------------");

    int selected = random(0, 4);

    Serial.print("Round "); Serial.print(currentRound + 1);
    Serial.print(" - LED "); Serial.print(selected + 1); Serial.println(" ON");

    displayRound(currentRound + 1, selected + 1);

    // Reset 4017
    digitalWrite(RESET_PIN, HIGH);
    delay(10);
    digitalWrite(RESET_PIN, LOW);

    // Advance clock to selected LED
    for (int i = 0; i <= ledSteps[selected]; i++) {
      digitalWrite(CLOCK_PIN, HIGH);
      delay(50);
      digitalWrite(CLOCK_PIN, LOW);
      delay(50);
    }

    Serial.println("Waiting for button press...");
    unsigned long startTime = millis();
    while (digitalRead(BUTTON_INPUT_PIN) == LOW) {
      // wait for button press
    }
    unsigned long responseTime = millis() - startTime;
    reactionTimes[currentRound] = responseTime;

    Serial.print("Button input received. Time: ");
    Serial.print(responseTime); Serial.println(" ms");

    delay(1000); // short pause between rounds
  }
}

void waitForTouchToContinue() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.println("Touch sensor to");
  display.println("continue...");
  display.println("---------------------");
  display.display();

  Serial.println("Waiting for touch sensor to continue...");
  while (digitalRead(TOUCH_SENSOR_PIN) == LOW) {
    // wait
  }
  Serial.println("Touch detected, continuing...");

  delay(300); // debounce
  runCountdown();
}

void displayRound(int roundNum, int ledNum) {
  display.clearDisplay();
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.print("Round "); display.println(roundNum);
  display.print("LED "); display.print(ledNum); display.println(" ON");
  display.display();
  Serial.print("Displaying Round "); Serial.print(roundNum);
  Serial.print(" | LED "); Serial.print(ledNum);
  Serial.println(" ON");
}

void showResults() {
  unsigned long total = 0;
  for (int i = 0; i < NUM_ROUNDS; i++) {
    total += reactionTimes[i];
  }
  unsigned long avg = total / NUM_ROUNDS;

  display.clearDisplay();
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.println("Game Over");
  display.print("Average: "); display.print(avg); display.println(" ms");
  display.display();

  Serial.println("---------------------");
  Serial.println("Game Over!");
  for (int i = 0; i < NUM_ROUNDS; i++) {
    Serial.print("Round "); Serial.print(i + 1);
    Serial.print(" time: "); Serial.print(reactionTimes[i]); Serial.println(" ms");
  }
  Serial.print("Average Time: "); Serial.print(avg); Serial.println(" ms");
  Serial.println("---------------------");
}
