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

#define PIN_RED   17
#define PIN_GREEN 16
#define PIN_BLUE  25

#define N_TOUCH   6
#define THRESHOLD 30

// Touch pin order: P03,P04,P02,P27,P01,P26
int  touch_pins[N_TOUCH]   = {3, 4, 2, 27, 1, 26};
int  touch_values[N_TOUCH] = {0, 0, 0, 0, 0, 0};
bool pin_touched_now[N_TOUCH]  = {false, false, false, false, false, false};
bool pin_touched_past[N_TOUCH] = {false, false, false, false, false, false};

// Arrow mappings (indices into touch_pins)
#define Q2_IDX 2  // DOWN  -> P02
#define Q3_IDX 3  // LEFT  -> P27
#define Q4_IDX 4  // RIGHT -> P01
#define Q5_IDX 5  // UP    -> P26

// Paint/erase pads
#define PAINT_WHITE_IDX 0  // P03 -> set cell white
#define PAINT_BLACK_IDX 1  // P04 -> set cell black

// ---- Display (0.96" OLED, typical SSD1306 128x64) ----
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_ADDR_PRIMARY 0x3C
#define OLED_ADDR_ALT     0x3D
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// Pixel-art grid: 16x8 of 8x8 cells
#define CELL      8
#define GRID_COLS (SCREEN_WIDTH / CELL)   // 16
#define GRID_ROWS (SCREEN_HEIGHT / CELL)  // 8
bool grid[GRID_ROWS][GRID_COLS]; // false=black, true=white

// Cursor (which "pixel" is selected)
int curRow = 0;
int curCol = 0;

// Blink the selection
unsigned long lastBlink = 0;
const unsigned long BLINK_MS = 500;  // a bit slower for visibility on 0.96"
bool blinkOn = false;

// Optional: draw faint grid lines (off by default to keep it crisp)
const bool SHOW_GRID = false;

void setLED(bool r, bool g, bool b) {
  // Active-LOW LED pins
  digitalWrite(PIN_RED,   r ? LOW : HIGH);
  digitalWrite(PIN_GREEN, g ? LOW : HIGH);
  digitalWrite(PIN_BLUE,  b ? LOW : HIGH);
}

void update_touch() {
  int t;
  const int t_max = 200;
  int p;

  for (int i = 0; i < N_TOUCH; i++) {
    p = touch_pins[i];

    pinMode(p, OUTPUT);
    digitalWriteFast(p, LOW);
    delayMicroseconds(25);

    noInterrupts();
    pinMode(p, INPUT_PULLUP);

    t = 0;
    while (!digitalReadFast(p) && t < t_max) {
      t++;
    }
    touch_values[i] = t;
    interrupts();

    pin_touched_past[i] = pin_touched_now[i];
    pin_touched_now[i]  = (touch_values[i] > THRESHOLD);
  }
}

void print_touch() {
  char print_buffer[30];
  for (int i=0; i < N_TOUCH; i++) {
    sprintf(print_buffer, "%4d ", touch_values[i]);
    Serial.print(print_buffer);
  }
  Serial.println("");
}

void drawGrid(bool blinkBorder) {
  display.clearDisplay();

  // Optional grid lines
  if (SHOW_GRID) {
    for (int x = 0; x <= SCREEN_WIDTH; x += CELL) display.drawFastVLine(x, 0, SCREEN_HEIGHT, WHITE);
    for (int y = 0; y <= SCREEN_HEIGHT; y += CELL) display.drawFastHLine(0, y, SCREEN_WIDTH, WHITE);
  }

  // Draw filled cells
  for (int r = 0; r < GRID_ROWS; r++) {
    for (int c = 0; c < GRID_COLS; c++) {
      if (grid[r][c]) {
        // Fill whole cell white
        display.fillRect(c*CELL, r*CELL, CELL, CELL, WHITE);
      }
    }
  }

  // Highlight current cell with a blinking border.
  // Using INVERSE makes the border visible whether the cell is black or white.
  if (blinkBorder) {
    int x = curCol * CELL;
    int y = curRow * CELL;
    display.drawRect(x, y, CELL, CELL, INVERSE);
  }

  display.display();
}

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

  pinMode(PIN_RED, OUTPUT);
  pinMode(PIN_GREEN, OUTPUT);
  pinMode(PIN_BLUE, OUTPUT);
  setLED(false, false, false); // off

  // I2C pins per your mapping (P06 SDA, P07 SCL), faster clock for snappier refresh
  Wire.setSDA(6);
  Wire.setSCL(7);
  Wire.begin();
  Wire.setClock(400000);

  // OLED init: try 0x3C then 0x3D
  bool ok = display.begin(SSD1306_SWITCHCAPVCC, OLED_ADDR_PRIMARY);
  if (!ok) {
    ok = display.begin(SSD1306_SWITCHCAPVCC, OLED_ADDR_ALT);
  }
  if (ok) {
    display.clearDisplay();
    // If your module is physically upside down, uncomment:
    // display.setRotation(2);
    display.display();
  }

  // Clear grid
  for (int r = 0; r < GRID_ROWS; r++)
    for (int c = 0; c < GRID_COLS; c++)
      grid[r][c] = false;

  drawGrid(blinkOn);
}

void loop() {
  update_touch();

  // LED color while held (UP=Red, DOWN=Green, LEFT=Blue, RIGHT=Yellow)
  if      (pin_touched_now[Q5_IDX]) setLED(true,  false, false); // Red
  else if (pin_touched_now[Q2_IDX]) setLED(false, true,  false); // Green
  else if (pin_touched_now[Q3_IDX]) setLED(false, false, true ); // Blue
  else if (pin_touched_now[Q4_IDX]) setLED(true,  true,  false); // Yellow
  else                              setLED(false, false, false); // off

  bool needRedraw = false;

  // Move cursor on PRESS (edge-triggered)
  if (pin_touched_now[Q5_IDX] && !pin_touched_past[Q5_IDX]) {           // UP
    if (curRow > 0) curRow--;
    needRedraw = true;
  }
  if (pin_touched_now[Q2_IDX] && !pin_touched_past[Q2_IDX]) {           // DOWN
    if (curRow < GRID_ROWS - 1) curRow++;
    needRedraw = true;
  }
  if (pin_touched_now[Q3_IDX] && !pin_touched_past[Q3_IDX]) {           // LEFT
    if (curCol > 0) curCol--;
    needRedraw = true;
  }
  if (pin_touched_now[Q4_IDX] && !pin_touched_past[Q4_IDX]) {           // RIGHT
    if (curCol < GRID_COLS - 1) curCol++;
    needRedraw = true;
  }

  // Paint / Erase on PRESS
  if (pin_touched_now[PAINT_WHITE_IDX] && !pin_touched_past[PAINT_WHITE_IDX]) {
    grid[curRow][curCol] = true;   // white
    needRedraw = true;
  }
  if (pin_touched_now[PAINT_BLACK_IDX] && !pin_touched_past[PAINT_BLACK_IDX]) {
    grid[curRow][curCol] = false;  // black/off
    needRedraw = true;
  }

  // Blink timing
  unsigned long now = millis();
  if (now - lastBlink >= BLINK_MS) {
    lastBlink = now;
    blinkOn = !blinkOn;
    needRedraw = true;
  }

  if (needRedraw) {
    drawGrid(blinkOn);
  }

  print_touch();
  delay(30);
}