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

#define SCREEN_WIDTH 128  // OLED display width, in pixels
#define SCREEN_HEIGHT 64  // OLED display height, in pixels
#define BLOCK_SIZE 4      // Group cells in 4x4 blocks
#define GRID_WIDTH (SCREEN_WIDTH / BLOCK_SIZE)  // Effective grid width
#define GRID_HEIGHT (SCREEN_HEIGHT / BLOCK_SIZE) // Effective grid height
#define OLED_RESET -1     // Reset pin (or -1 if sharing Arduino reset pin)
#define SCREEN_ADDRESS 0x3C  // Most common I2C address for 128x64 OLED displays

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

int grid[GRID_WIDTH][GRID_HEIGHT];
int nextGrid[GRID_WIDTH][GRID_HEIGHT];

void setup() {
  // Initialize the display
  if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);  // Don't proceed, loop forever
  }
  display.clearDisplay();

  initializeRPentomino();  // Initialize the grid with the R-pentomino
  display.display();
}

void loop() {
  displayGrid();  // Display the current grid state
  updateGrid();   // Update the grid for the next generation
  delay(100);     // Adjust delay for visual effect
}

// Initialize the grid with the R-pentomino pattern in the center
void initializeRPentomino() {
  // Clear the grid
  for (int x = 0; x < GRID_WIDTH; x++) {
    for (int y = 0; y < GRID_HEIGHT; y++) {
      grid[x][y] = 0;  // All cells dead initially
    }
  }

  // Define R-pentomino in the center of the grid
  int centerX = GRID_WIDTH / 2;
  int centerY = GRID_HEIGHT / 2;

  // R-pentomino pattern
  grid[centerX][centerY] = 1;
  grid[centerX + 1][centerY] = 1;
  grid[centerX - 1][centerY + 1] = 1;
  grid[centerX][centerY + 1] = 1;
  grid[centerX][centerY + 2] = 1;
}

// Display the grid on the OLED
void displayGrid() {
  display.clearDisplay();
  for (int x = 0; x < GRID_WIDTH; x++) {
    for (int y = 0; y < GRID_HEIGHT; y++) {
      if (grid[x][y] == 1) {
        // Draw 4x4 block for each live cell
        display.fillRect(x * BLOCK_SIZE, y * BLOCK_SIZE, BLOCK_SIZE, BLOCK_SIZE, SSD1306_WHITE);
      }
    }
  }
  display.display();  // Refresh the display
}

// Count live neighbors around a cell
int countNeighbors(int x, int y) {
  int count = 0;
  for (int i = -1; i <= 1; i++) {
    for (int j = -1; j <= 1; j++) {
      if (i == 0 && j == 0) continue;  // Skip the cell itself
      int newX = (x + i + GRID_WIDTH) % GRID_WIDTH;  // Wrap around edges
      int newY = (y + j + GRID_HEIGHT) % GRID_HEIGHT;
      count += grid[newX][newY];  // Count live neighbors
    }
  }
  return count;
}

// Update the grid based on Game of Life rules
void updateGrid() {
  for (int x = 0; x < GRID_WIDTH; x++) {
    for (int y = 0; y < GRID_HEIGHT; y++) {
      int neighbors = countNeighbors(x, y);
      if (grid[x][y] == 1) {
        if (neighbors < 2 || neighbors > 3) {
          nextGrid[x][y] = 0;  // Cell dies due to under/overpopulation
        } else {
          nextGrid[x][y] = 1;  // Cell survives
        }
      } else {
        if (neighbors == 3) {
          nextGrid[x][y] = 1;  // Cell becomes alive
        } else {
          nextGrid[x][y] = 0;  // Cell stays dead
        }
      }
    }
  }

  // Copy nextGrid to grid for the next generation
  for (int x = 0; x < GRID_WIDTH; x++) {
    for (int y = 0; y < GRID_HEIGHT; y++) {
      grid[x][y] = nextGrid[x][y];
    }
  }
}