/*
  Traffic Light with Crosswalk Display for Uno
  AnonEngineering
  7/14/25
  Beerware
*/

#include <Wire.h>
#include <U8g2lib.h>
#include <timeObj.h>  // from LC_baseTools
#include "bitmaps.h"

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET    -1

typedef enum { STATE_INIT,
               STATE_NS_GO,
               STATE_NS_COUNT,
               STATE_NS_CAUTION,
               STATE_NS_ALL_STOP,
               STATE_EW_GO,
               STATE_EW_COUNT,
               STATE_EW_CAUTION,
               STATE_EW_ALL_STOP
             } states;

// set initial state
states state = STATE_INIT;
states oldState = state;
states nextState = STATE_NS_GO;

// sets the number of seconds in each state or the countdown count
const int STATE_DELAYS[] = {0, 20, 10, 3, 2, 20, 10, 3, 2};
const char STATE_NAME[][16] = {
  {"Initialized"},
  {"NS go"}, {"NS count"}, {"NS caution"}, {"NS all stop"},
  {"EW go"}, {"EW count"}, {"EW caution"}, {"EW all stop"}
};
const unsigned long ONE_SEC = 1000;
const int NUM_BTNS = 2;
const int BTN_PINS[] = {9, 5};
const int NS_LED_PINS[] = {12, 11, 10};
const int EW_LED_PINS[] = {8, 7, 6};

int btnState[NUM_BTNS];
int oldBtnState[NUM_BTNS];

U8G2_SSD1306_128X64_NONAME_1_HW_I2C u8g2_ns(U8G2_R2, U8X8_PIN_NONE);
U8G2_SSD1306_128X64_NONAME_1_HW_I2C u8g2_ew(U8G2_R2, U8X8_PIN_NONE);
timeObj  stateTimer;
timeObj  cdTimer;

// function returns which button was pressed, or 0 if none
int checkButtons()  {
  int btnPressed = 0;

  for (int i = 0; i < NUM_BTNS; i++) {
    btnState[i] = digitalRead(BTN_PINS[i]); // check each button
    if (btnState[i] != oldBtnState[i])  {   // if it changed
      oldBtnState[i] = btnState[i];         // remember state for next time
      if (btnState[i] == 0)  {              // was just pressed
        btnPressed = i + 1;
        if (btnPressed == 1)  {
          Serial.print("North/South ");
        }
        if (btnPressed == 2)  {
          Serial.print("East/West ");
        }
        Serial.println("crosswalk button pressed");
      }
      delay(20);  // debounce
    }
  }
  return btnPressed;
}

void initialize() {
  u8g2_ns.setDrawColor(1);
  u8g2_ns.firstPage();
  do {
    u8g2_ns.drawStr(0, 20, "North");
  } while ( u8g2_ns.nextPage() );
  u8g2_ew.setDrawColor(1);
  u8g2_ew.firstPage();
  do {
    u8g2_ew.drawStr(0, 20, "East");
  } while ( u8g2_ew.nextPage() );
  delay(2000);
  int lit[] = {0, 0, 1, 1, 0, 0};
  for (int pin = 0; pin < 3; pin++) {
    digitalWrite(NS_LED_PINS[pin], lit[pin]);
    digitalWrite(EW_LED_PINS[pin], lit[pin + 3]);
  }
  // clear the buffers
  u8g2_ns.clear();
  u8g2_ew.clear();

  u8g2_ns.setDrawColor(0);
  u8g2_ns.firstPage();
  do {
    u8g2_ns.drawXBMP(32, 0, 64, 64, oled_64x64_images[0]);
  } while ( u8g2_ns.nextPage() );

  u8g2_ew.setDrawColor(0);
  u8g2_ew.firstPage();
  do {
    u8g2_ew.drawXBMP(16, 0, 64, 64, oled_64x64_images[1]);
  } while ( u8g2_ew.nextPage() );

  printState();
  state = STATE_NS_GO;
  stateTimer.setTime(STATE_DELAYS[1] * 1000);
  stateTimer.start();
}

void nsGo() {
  int lit[] = {0, 0, 1, 1, 0, 0};
  for (int pin = 0; pin < 3; pin++) {
    digitalWrite(NS_LED_PINS[pin], lit[pin]);
    digitalWrite(EW_LED_PINS[pin], lit[pin + 3]);
  }
  u8g2_ns.setDrawColor(0);
  u8g2_ns.firstPage();
  do {
    u8g2_ns.drawXBMP(32, 0, 64, 64, oled_64x64_images[0]);
  } while ( u8g2_ns.nextPage() );
}

bool nsCount()  {
  static int count = STATE_DELAYS[2] + 1;
  char buffer[3];

  if (cdTimer.ding())  {  ////// never started?
    cdTimer.setTime(ONE_SEC);
    cdTimer.start();
    count--;
    if (count < 0)  {
      count = STATE_DELAYS[2] + 1;
      return true;
    }
    snprintf(buffer, 3, "%2d", count);
    u8g2_ns.firstPage();
    do {
      u8g2_ns.setDrawColor(0);
      u8g2_ns.drawXBMP(16, 0, 64, 64, oled_64x64_images[1]);
      u8g2_ns.setDrawColor(1);
      u8g2_ns.drawStr(80, 40, buffer);
    } while ( u8g2_ns.nextPage() );
  }
  return false;
}

void nsCaution()  {
  int lit[] = {0, 1, 0, 1, 0, 0};
  for (int pin = 0; pin < 3; pin++) {
    digitalWrite(NS_LED_PINS[pin], lit[pin]);
    digitalWrite(EW_LED_PINS[pin], lit[pin + 3]);
  }
  u8g2_ns.setDrawColor(0);
  u8g2_ns.firstPage();
  do {
    u8g2_ns.drawXBMP(16, 0, 64, 64, oled_64x64_images[1]);
  } while ( u8g2_ns.nextPage() );
}

void nsAllStop()  {
  int lit[] = {1, 0, 0, 1, 0, 0};
  for (int pin = 0; pin < 3; pin++) {
    digitalWrite(NS_LED_PINS[pin], lit[pin]);
    digitalWrite(EW_LED_PINS[pin], lit[pin + 3]);
  }
}

void ewGo() {
  int lit[] = {1, 0, 0, 0, 0, 1};
  for (int pin = 0; pin < 3; pin++) {
    digitalWrite(NS_LED_PINS[pin], lit[pin]);
    digitalWrite(EW_LED_PINS[pin], lit[pin + 3]);
  }
  u8g2_ew.setDrawColor(0);
  u8g2_ew.firstPage();
  do {
    u8g2_ew.drawXBMP(32, 0, 64, 64, oled_64x64_images[0]);
  } while ( u8g2_ew.nextPage() );
}

bool ewCount()  {
  static int count = STATE_DELAYS[6] + 1;
  char buffer[3];

  if (cdTimer.ding())  {  ////// never started?
    cdTimer.setTime(ONE_SEC);
    cdTimer.start();
    count--;
    if (count < 0)  {
      count = STATE_DELAYS[6] + 1;
      return true;
    }
    snprintf(buffer, 3, "%2d", count);
    u8g2_ew.firstPage();
    do {
      u8g2_ew.setDrawColor(0);
      u8g2_ew.drawXBMP(16, 0, 64, 64, oled_64x64_images[1]);
      u8g2_ew.setDrawColor(1);
      u8g2_ew.drawStr(80, 40, buffer);
    } while ( u8g2_ew.nextPage() );
  }
  return false;
}

void ewCaution()  {
  int lit[] = {1, 0, 0, 0, 1, 0};
  for (int pin = 0; pin < 3; pin++) {
    digitalWrite(NS_LED_PINS[pin], lit[pin]);
    digitalWrite(EW_LED_PINS[pin], lit[pin + 3]);
  }
  u8g2_ew.setDrawColor(0);
  u8g2_ew.firstPage();
  do {
    u8g2_ew.drawXBMP(16, 0, 64, 64, oled_64x64_images[1]);
  } while ( u8g2_ew.nextPage() );
}

void ewAllStop()  {
  int lit[] = {1, 0, 0, 1, 0, 0};
  for (int pin = 0; pin < 3; pin++) {
    digitalWrite(NS_LED_PINS[pin], lit[pin]);
    digitalWrite(EW_LED_PINS[pin], lit[pin + 3]);
  }
}

void printState() {
  Serial.println(STATE_NAME[state]);
}

void setState() {
  bool isDone = false;

  if (state != oldState)  {
    oldState = state;
    printState();
  }

  switch (state) {
    case STATE_INIT:
      initialize();
      break;
    case STATE_NS_GO:
      nsGo();
      nextState = STATE_NS_COUNT;
      if (stateTimer.ding())  {
        state = nextState;
      }
      break;
    case STATE_NS_COUNT:
      isDone = nsCount();
      nextState = STATE_NS_CAUTION;
      if (isDone)  {
        state = nextState;
        stateTimer.setTime(STATE_DELAYS[3] * 1000);
        stateTimer.start();
      }
      break;
    case STATE_NS_CAUTION:
      nsCaution();
      nextState = STATE_NS_ALL_STOP;
      if (stateTimer.ding())  {
        state = nextState;
        stateTimer.setTime(STATE_DELAYS[4] * 1000);
        stateTimer.start();
      }
      break;
    case STATE_NS_ALL_STOP:
      nsAllStop();
      nextState = STATE_EW_GO;
      if (stateTimer.ding())  {
        state = nextState;
        stateTimer.setTime(STATE_DELAYS[5] * 1000);
        stateTimer.start();
      }
      break;
    case STATE_EW_GO:
      ewGo();
      nextState = STATE_EW_COUNT;
      if (stateTimer.ding())  {
        state = nextState;
      }
      break;
    case STATE_EW_COUNT:
      isDone = ewCount();
      nextState = STATE_EW_CAUTION;
      if (isDone)  {
        state = nextState;
        stateTimer.setTime(STATE_DELAYS[7] * 1000);
        stateTimer.start();
      }
      break;
    case STATE_EW_CAUTION:
      ewCaution();
      nextState = STATE_EW_ALL_STOP;
      if (stateTimer.ding())  {
        state = nextState;
        stateTimer.setTime(STATE_DELAYS[8] * 1000);
        stateTimer.start();
      }
      break;
    case STATE_EW_ALL_STOP:
      ewAllStop();
      nextState = STATE_NS_GO;
      if (stateTimer.ding())  {
        state = nextState;
        stateTimer.setTime(STATE_DELAYS[1] * 1000);
        stateTimer.start();
      }
      break;
    default:
      Serial.println("WTF");
  }
}

void setup() {
  Serial.begin(9600);
  for (int i = 0; i < 3; i++) {
    pinMode(NS_LED_PINS[i], OUTPUT);
    pinMode(EW_LED_PINS[i], OUTPUT);
  }
  pinMode(BTN_PINS[0], INPUT_PULLUP);
  pinMode(BTN_PINS[1], INPUT_PULLUP);
  u8g2_ns.setI2CAddress(0x3C * 2);  // 8 bit address
  u8g2_ew.setI2CAddress(0x3D * 2);
  u8g2_ns.begin();
  u8g2_ew.begin();
  u8g2_ns.setFont(u8g2_font_spleen16x32_mf);
  u8g2_ew.setFont(u8g2_font_spleen16x32_mf);
  Serial.println();
}

void loop() {
  int btnNumber = checkButtons();
  if (btnNumber)  {
    if (btnNumber == 1 && state == STATE_EW_GO)  {
      state = STATE_EW_COUNT;
    }
    if (btnNumber == 2 && state == STATE_NS_GO)  {
      state = STATE_NS_COUNT;
    }
  }
  setState();
}