#include <Wire.h>
#include <LiquidCrystal_I2C.h>

// -------------------- CONFIG: pins & constants --------------------

// I2C LCD
#define LCD_I2C_ADDR 0x27
LiquidCrystal_I2C lcd(LCD_I2C_ADDR, 20, 4);

// Debounce and timing
const unsigned long DEBOUNCE_MS = 40;
const unsigned long CONTROL_SHORTPRESS_MAX_MS = 500;

unsigned long lastDebounceTimeControl = 0;
bool lastControlRaw = false;
bool lastControlStable = false;
unsigned long controlDownTime = 0;

// LED pin arrays (change these to match your wiring)
const uint8_t greenPins[12] = {
  22, 24, 26, 28, 30, 32, 34, 36, 38, 40, 42, 44
};
const uint8_t bluePins[8] = {  // N, NE, NW, S, SE, SW, E, W, NULL
  A0, A1, A2, A3, A4, A5, A6, A7
};
const uint8_t yellowPins[8] = {
  A8, A9, A10, A11, A12, A13, A14, A15
};
int orangeArduinoIndex = 3;      // Set orange led index (this will be first orange led pin number)
const uint8_t orangePins[4] = {  // E, W, N, S
  3, 4, 5, 6
};

// Pushbuttons
const uint8_t controlButtonPin = 12;  // short press = turn off active LEDs
const uint8_t pbPins[15] = {
  23, 25, 27, 29, 31, 33, 35, 37, 39, 41, 43, 45, 47, 49, 51
};

// Toggle switches (each toggle uses two pins: UP and DOWN)
// Toggle A selects N / OFF / S  (UP => N, DOWN => S)
// Toggle B selects E / OFF / W  (UP => E, DOWN => W)
const uint8_t toggleA_up_pin = 8;
const uint8_t toggleA_down_pin = 9;
const uint8_t toggleB_up_pin = 10;
const uint8_t toggleB_down_pin = 11;

// LED logical states
bool greenState[12];
bool blueState[8];
bool yellowState[8];
bool orangeState[4];

// Current direction index (0 = NULL, 1=N,2=S,3=E,4=W,5=NE,6=NW,7=SE,8=SW)
uint8_t currentDirectionIndex = 15;

uint8_t currentScenario = 1;  // set externally if you have scenario-select hardware
bool anyLedOn = false;        // FLag to check if any pb was pressed previously and led is on

// ---------- Compact PB -> Direction -> Green Indices array ----------
// pbMap[PB_index][directionIndex][slot]  (max 4 entries per mapping, -1 terminator)
// Direction order follows constants: 0=NULL,1=N,2=S,3=E,4=W,5=NE,6=NW,7=SE,8=SW

const int8_t pbMap[15][9][3] = {
  // PB1
  { { 4, 0, -1 }, { 4, 0, -1 }, { 6, 1, -1 }, { 5, 6, 3 }, { 5, 6, 3 }, { 4, 5, 3 }, { 4, 5, 0 }, { 6, 3, -1 }, { 6, 3, -1 } },
  // PB2
  { { 0, 11, 2 }, { 0, 11, 2 }, { 0, 11, 2 }, { 1, 0, -1 }, { 11, 2, -1 }, { 1, 0, -1 }, { 1, 0, -1 }, { 11, 2, -1 }, { 11, 2, -1 } },
  // PB3
  { { 7, 1, -1 }, { 7, 1, -1 }, { 5, 3, -1 }, { 6, 3, -1 }, { 6, 7, 3 }, { 6, 3, -1 }, { 5, 3, -1 }, { 7, 1, -1 }, { 7, 1, -1 } },
  // PB4
  { { 3, 2, 0 }, { 3, 2, 0 }, { 3, 2, 0 }, { 4, 3, -1 }, { 4, 3, -1 }, { 3, 1, -1 }, { 3, 2, 1 }, { 4, 3, -1 }, { 2, 0, -1 } },
  // PB5
  { { 2, 0, -1 }, { 2, 0, -1 }, { 3, 2, 0 }, { 3, 3, -1 }, { 3, 3, -1 }, { 3, 3, -1 }, { 3, 2, 0 }, { 2, 2, -1 }, { 2, 2, -1 } },
  // PB6
  { { 2, 2, -1 }, { 2, 2, -1 }, { 3, 2, 0 }, { 3, 0, -1 }, { 3, 0, -1 }, { 3, 2, 0 }, { 2, 0, -1 }, { 3, 3, -1 }, { 1, 2, -1 } },
  // PB7
  { { 1, 2, -1 }, { 1, 2, -1 }, { 1, 2, 0 }, { 2, 0, -1 }, { 2, 0, -1 }, { 3, 2, 0 }, { 2, 1, 0 }, { 1, 2, -1 }, { 3, 0, -1 } },
  // PB8
  { { 1, 2, -1 }, { 1, 2, -1 }, { 1, 2, -1 }, { 2, 0, -1 }, { 2, 0, -1 }, { 2, 1, 0 }, { 2, 1, 0 }, { 2, 0, -1 }, { 2, 0, -1 } },
  // PB9
  { { 8, 1, -1 }, { 7, 1, -1 }, { 8, 1, -1 }, { 6, 7, 3 }, { 6, 7, 3 }, { 6, 7, 3 }, { 8, 1, -1 }, { 7, 8, 1 }, { 8, 1, -1 } },
  // PB10
  { { 8, 1, -1 }, { 8, 1, -1 }, { 8, 1, -1 }, { 7, 3, -1 }, { 7, 3, -1 }, { 6, 7, 3 }, { 7, 3, -1 }, { 7, 8, 1 }, { 8, 1, -1 } },
  // PB11
  { { 9, 1, -1 }, { 9, 1, -1 }, { 9, 1, -1 }, { 8, 3, -1 }, { 8, 1, -1 }, { 6, 7, 3 }, { 7, 3, -1 }, { 8, 9, 1 }, { 7, 3, -1 } },
  // PB12
  { { 9, 2, -1 }, { 9, 1, -1 }, { 9, 2, -1 }, { 8, 1, -1 }, { 8, 9, 1 }, { 8, 1, -1 }, { 8, 3, -1 }, { 8, 9, 1 }, { 7, 8, 3 } },
  // PB13
  { { 10, 2, -1 }, { 9, 10, 2 }, { 10, 2, -1 }, { 9, 1, -1 }, { 9, 1, -1 }, { 9, 1, -1 }, { 10, 2, -1 }, { 9, 10, 1 }, { 8, 1, -1 } },
  // PB14
  { { 10, 2, -1 }, { 10, 2, -1 }, { 10, 2, -1 }, { 9, 1, -1 }, { 9, 10, 1 }, { 9, 1, -1 }, { 9, 1, -1 }, { 9, 10, 1 }, { 9, 1, -1 } },
  // PB15
  { { 11, 2, -1 }, { 11, 2, -1 }, { 11, 2, -1 }, { 10, 1, -1 }, { 9, 10, 1 }, { 9, 1, -1 }, { 11, 2, -1 }, { 11, 1, -1 }, { 9, 1, -1 } }
};

void applyBluePattern(uint8_t dirIdx) {
  setLedState(bluePins, blueState, 8, LOW, HIGH);
  // else
  digitalWrite(bluePins[dirIdx], HIGH);  // Turn on blue led according to index
}
// Read toggle that has up/down pins. Return +1 for UP, 0 for center, -1 for DOWN.
int8_t read3PosToggle(uint8_t upPin, uint8_t downPin) {
  bool up = digitalRead(upPin);
  bool down = digitalRead(downPin);
  if (up && !down) return 1;
  if (down && !up) return -1;
  return 0;
}

// Convert toggle combination to direction index
uint8_t getDirectionFromToggles(int8_t tA, int8_t tB) {
  // tA: -1=N, 0=OFF, 1=S
  // tB: -1=E, 0=OFF, 1=W
  // Serial.print("ta:");
  // Serial.print(tA);
  // Serial.print(", tB:");
  // Serial.println(tB);
  if (tA == -1 && tB == 0) return 0;   // Return North
  if (tA == -1 && tB == -1) return 1;  // Return NE
  if (tA == -1 && tB == 1) return 2;   // Return NW
  if (tA == 1 && tB == 0) return 3;    // Return S
  if (tA == 1 && tB == -1) return 4;   // Return SE
  if (tA == 1 && tB == 1) return 5;    // Return SW
  if (tA == 0 && tB == -1) return 6;   // Return E
  if (tA == 0 && tB == 1) return 7;    // Return W
  return 8;                            // Return NULL
}
// Helper to get direction string for LCD display
const char* directionNames[9] = {
  "N", "NE", "NW", "S", "SE", "SW", "E", "W", "--"
};

// -------------------- SETUP --------------------
void setupPins() {
  setLedState(greenPins, greenState, 12, LOW, HIGH);
  setLedState(bluePins, blueState, 8, LOW, HIGH);
  setLedState(yellowPins, yellowState, 8, LOW, HIGH);
  setLedState(orangePins, orangeState, 4, LOW, HIGH);


  // Buttons
  pinMode(controlButtonPin, INPUT_PULLUP);  // active LOW (pressed => LOW)
  for (int i = 0; i < 15; i++) {
    pinMode(pbPins[i], INPUT_PULLUP);  // active LOW
  }
  pinMode(toggleA_up_pin, INPUT_PULLUP);
  pinMode(toggleA_down_pin, INPUT_PULLUP);
  pinMode(toggleB_up_pin, INPUT_PULLUP);
  pinMode(toggleB_down_pin, INPUT_PULLUP);
}
void setLedState(const uint8_t pins[], bool states[], size_t count, bool ledState, bool initialize) {
  for (size_t i = 0; i < count; i++) {
    if (initialize) pinMode(pins[i], OUTPUT);
    digitalWrite(pins[i], ledState);  // Your custom safe function
    states[i] = false;
  }
}
void lcdPrintAt(int row, int col, const char* text, bool clear) {
  if (clear) lcd.clear();  // If clear flag was set clear lcd
  lcd.setCursor(col, row);
  lcd.print(text);
  lcd.print("  ");
}
void lcdPrintAt(int row, int col, uint8_t value, bool clear) {
  if (clear) lcd.clear();  // If clear flag was set clear lcd
  lcd.setCursor(col, row);
  lcd.print(value);
  lcd.print("  ");
}
void setup() {
  Serial.begin(115200);
  setupPins();
  // Init LCD
  lcd.init();
  lcd.backlight();
  lcdPrintAt(0, 0, "LED Control", HIGH);
  lcdPrintAt(1, 0, "Arduino Mega", LOW);
  setupPins();
  delay(800);
  printScenarioOnDisplay();
}

void loop() {
  readToggleSwitches();  // Read toggle switches and take action

  bool rawControl = (digitalRead(controlButtonPin) == LOW);  // pressed active LOW
  unsigned long now = millis();
  if (rawControl != lastControlRaw) {
    lastDebounceTimeControl = now;
    lastControlRaw = rawControl;
  }
  if ((now - lastDebounceTimeControl) > DEBOUNCE_MS) {
    if (rawControl != lastControlStable) {
      lastControlStable = rawControl;
      if (lastControlStable) {
        // button just pressed
        controlDownTime = now;
      } else {
        // button released: check press length
        unsigned long len = now - controlDownTime;
        Serial.print("Control button was pressed for ");
        Serial.println(len);
        if (len <= CONTROL_SHORTPRESS_MAX_MS) {
          Serial.println("Control Button Short Pressed");
          if (anyLedOn) {                                       // If previously led is on turn off them
            anyLedOn = false;                                   // Reset the flag
            setLedState(greenPins, greenState, 12, LOW, LOW);   // Turn off green leds
            setLedState(orangePins, orangeState, 4, LOW, LOW);  // Turn off orange leds
          }
          // short press: turn off all PB-activated LEDs (not blue)
          // handleDirectionChangeReset();
        } else if (len > 3000) {
          currentScenario = currentScenario + 1 > 7 ? 1 : currentScenario + 1;
          if (currentScenario > 2) {
            lcdPrintAt(0, 1, "Shelter in Place", LOW);
            setLedState(yellowPins, yellowState, 8, HIGH, LOW);  // Turn on yellow leds
          } else {
            setLedState(yellowPins, yellowState, 8, LOW, LOW);  // Turn off yellow leds
            lcdPrintAt(0, 0, "Safe            ", LOW);
          }
          printScenarioOnDisplay();

          if (anyLedOn) {                                       // If previously led is on turn off them
            anyLedOn = false;                                   // Reset the flag
            setLedState(greenPins, greenState, 12, LOW, LOW);   // Turn off green leds
            setLedState(orangePins, orangeState, 4, LOW, LOW);  // Turn off orange leds
          }
          Serial.print("Button pressed for more than 3 second, change scenario to:");
          Serial.println(currentScenario);
          // long press (not specified) - ignore or implement later
        }
      }
    }
  }
  if (currentScenario <= 2) {  // Check if we are on scenarion 1 and 2
    for (int i = 0; i < 15; i++) {
      if (digitalRead(pbPins[i]) == LOW) {                    // Check if any button is pressed
        if (anyLedOn) {                                       // If previously led is on turn off them
          anyLedOn = false;                                   // Reset the flag
          setLedState(greenPins, greenState, 12, LOW, LOW);   // Turn off green leds
          setLedState(orangePins, orangeState, 4, LOW, LOW);  // Turn off orange leds
        } else {                                              // If previously led is off turn it on
          anyLedOn = true;                                    // Set the flag

          // Check if the 3rd (index 2) element is not -1 then first two elements will belong to Green led and 3rd element will be for orange led
          // If 3rd (index 2) element is -1 then index 0 element will be for green led and index 1 element will be for orange led
          // Serial.print("Button index:");
          // Serial.print(i);
          // Serial.print(", dir index:");
          // Serial.print(currentDirectionIndex);
          // Serial.print(", pbMap 0:");
          // Serial.print(pbMap[i][currentDirectionIndex][0]);
          // Serial.print(", pbMap 1:");
          // Serial.print(pbMap[i][currentDirectionIndex][1]);
          // Serial.print(", pbMap 2:");
          // Serial.print(pbMap[i][currentDirectionIndex][2]);

          if (pbMap[i][currentDirectionIndex][2] == -1) {
            digitalWrite(greenPins[pbMap[i][currentDirectionIndex][0]], HIGH);
            digitalWrite(orangePins[pbMap[i][currentDirectionIndex][1]], HIGH);
            // Serial.print(", green index 1:");
            // Serial.print(greenPins[pbMap[i][currentDirectionIndex][0]]);
            // Serial.print(", orange index 2:");
            // Serial.println(orangePins[pbMap[i][currentDirectionIndex][1]]);
          } else {  // First 2 indexes will be for green, and 3rd index will be for orange
            digitalWrite(greenPins[pbMap[i][currentDirectionIndex][0]], HIGH);
            digitalWrite(greenPins[pbMap[i][currentDirectionIndex][1]], HIGH);
            digitalWrite(orangePins[pbMap[i][currentDirectionIndex][2]], HIGH);
            // Serial.print(", green index 1:");
            // Serial.print(greenPins[pbMap[i][currentDirectionIndex][0]]);
            // Serial.print(", green index 2:");
            // Serial.print(greenPins[pbMap[i][currentDirectionIndex][1]]);
            // Serial.print(", ornage index 3:");
            // Serial.println(orangePins[pbMap[i][currentDirectionIndex][2]]);
          }
        }
        delay(200);  // Debounce delay
      }
    }
  }
}
void printScenarioOnDisplay() {
  lcdPrintAt(1, 13, "S:", LOW);
  lcdPrintAt(1, 15, currentScenario, LOW);
}

void readToggleSwitches() {
  // 1) Read toggles and compute direction
  int8_t tA = read3PosToggle(toggleA_up_pin, toggleA_down_pin);
  int8_t tB = read3PosToggle(toggleB_up_pin, toggleB_down_pin);
  uint8_t newDir = getDirectionFromToggles(tA, tB);

  // If direction changes, turn off all leds
  if (newDir != currentDirectionIndex) {
    currentDirectionIndex = newDir;
    Serial.print("New direction index:");
    Serial.println(currentDirectionIndex);
    setLedState(greenPins, greenState, 12, LOW, LOW);  // Turn off green leds
    // setLedState(yellowPins, yellowState, 8, LOW, LOW);  // Turn off yellow leds
    for (int i = 0; i < 4; i++) {
      digitalWrite(orangePins[i], LOW);
    }
    lcdPrintAt(0, 14, directionNames[currentDirectionIndex], LOW);
    // setLedState(orangePins, orangeState, 4, LOW, LOW);  // Turn off orange leds
    // 2) Apply blue LEDs according to direction (blue LEDs always reflect toggle, independent of other state)
    applyBluePattern(currentDirectionIndex);
  }
}