const int redLED1 = 1;  // Red LED 1
const int yellowLED1 = 2;  // Yellow LED 1
const int greenLED1 = 3;  // Green LED 1
const int redLED2 = 4;  // Red LED 2
const int yellowLED2 = 5;  // Yellow LED 2
const int greenLED2 = 6;  // Green LED 2
const int redLED3 = 7;  // Red LED 3
const int yellowLED3 = 8;  // Yellow LED 3
const int greenLED3 = 9;  // Green LED 3
const int redLED4 = 10;  // Red LED 4
const int yellowLED4 = 11;  // Yellow LED 4
const int greenLED4 = 12;  // Green LED 4

const int button1 = 14;  // Push button for Traffic Light 1
const int button2 = 15;  // Push button for Traffic Light 2
const int button3 = 16;  // Push button for Traffic Light 3
const int button4 = 17;  // Push button for Traffic Light 4

bool buttonPressed = false;
unsigned long previousMillis = 0;
unsigned long interval = 1000; 
int currentStep = 0;
bool resumeFromYellow = false; 
int lastStep = -1;  // Track the last traffic light that was on

void setup() {
  pinMode(redLED1, OUTPUT);
  pinMode(yellowLED1, OUTPUT);
  pinMode(greenLED1, OUTPUT);
  pinMode(redLED2, OUTPUT);
  pinMode(yellowLED2, OUTPUT);
  pinMode(greenLED2, OUTPUT);
  pinMode(redLED3, OUTPUT);
  pinMode(yellowLED3, OUTPUT);
  pinMode(greenLED3, OUTPUT);
  pinMode(redLED4, OUTPUT);
  pinMode(yellowLED4, OUTPUT);
  pinMode(greenLED4, OUTPUT);

  pinMode(button1, INPUT_PULLUP);
  pinMode(button2, INPUT_PULLUP);
  pinMode(button3, INPUT_PULLUP);
  pinMode(button4, INPUT_PULLUP);
}

void loop() {
  if (digitalRead(button1) == LOW) {
    buttonPressed = true;
    interruptTraffic(greenLED1, redLED2, redLED3, redLED4, 0); // track the step
  } else if (digitalRead(button2) == LOW) {
    buttonPressed = true;
    interruptTraffic(greenLED2, redLED1, redLED3, redLED4, 1);
  } else if (digitalRead(button3) == LOW) {
    buttonPressed = true;
    interruptTraffic(greenLED3, redLED1, redLED2, redLED4, 2);
  } else if (digitalRead(button4) == LOW) {
    buttonPressed = true;
    interruptTraffic(greenLED4, redLED1, redLED2, redLED3, 3);
  } else {
    unsigned long currentMillis = millis();
    if (currentMillis - previousMillis >= interval) {
      previousMillis = currentMillis;
      
      if (resumeFromYellow) {
        resumeWithYellow();  // Resume from yellow
      } else {
        trafficLightSequence();  // Normal sequence
      }
    }
  }
}

void interruptTraffic(int greenPin, int otherRed1, int otherRed2, int otherRed3, int step) {
  turnAllOff();

  digitalWrite(otherRed1, HIGH);
  digitalWrite(otherRed2, HIGH);
  digitalWrite(otherRed3, HIGH);

  digitalWrite(greenPin, HIGH);

  lastStep = step;  // Remember which traffic light was active before interruption
  resumeFromYellow = true;

  while (digitalRead(button1) == LOW || digitalRead(button2) == LOW || digitalRead(button3) == LOW || digitalRead(button4) == LOW);
}

void turnAllOff() {
  digitalWrite(greenLED1, LOW);
  digitalWrite(yellowLED1, LOW);
  digitalWrite(redLED1, LOW);
  digitalWrite(greenLED2, LOW);
  digitalWrite(yellowLED2, LOW);
  digitalWrite(redLED2, LOW);
  digitalWrite(greenLED3, LOW);
  digitalWrite(yellowLED3, LOW);
  digitalWrite(redLED3, LOW);
  digitalWrite(greenLED4, LOW);
  digitalWrite(yellowLED4, LOW);
  digitalWrite(redLED4, LOW);
}

void trafficLightSequence() {
  switch (currentStep) {
    case 0:
      activateTrafficLight(greenLED1, yellowLED1, redLED1, redLED2, redLED3, redLED4);
      currentStep = 1;
      break;
    case 1:
      activateTrafficLight(greenLED2, yellowLED2, redLED2, redLED1, redLED3, redLED4);
      currentStep = 2;
      break;
    case 2:
      activateTrafficLight(greenLED3, yellowLED3, redLED3, redLED1, redLED2, redLED4);
      currentStep = 3;
      break;
    case 3:
      activateTrafficLight(greenLED4, yellowLED4, redLED4, redLED1, redLED2, redLED3);
      currentStep = 0;
      break;
  }
}

void activateTrafficLight(int greenPin, int yellowPin, int redPin, int otherRed1, int otherRed2, int otherRed3) {
  turnAllOff();

  digitalWrite(otherRed1, HIGH);
  digitalWrite(otherRed2, HIGH);
  digitalWrite(otherRed3, HIGH);

  digitalWrite(redPin, HIGH);

  unsigned long startTime = millis();
  while (millis() - startTime < interval) {
    if (checkButtonPress()) return;
  }

  digitalWrite(redPin, LOW);
  digitalWrite(yellowPin, HIGH);
  startTime = millis();
  while (millis() - startTime < interval) {
    if (checkButtonPress()) return;
  }

  digitalWrite(yellowPin, LOW);
  digitalWrite(greenPin, HIGH);
  startTime = millis();
  while (millis() - startTime < interval * 3) {
    if (checkButtonPress()) return;
  }

  digitalWrite(greenPin, LOW);
  digitalWrite(yellowPin, HIGH);
  startTime = millis();
  while (millis() - startTime < interval) {
    if (checkButtonPress()) return;
  }

  digitalWrite(yellowPin, LOW);
  digitalWrite(redPin, HIGH);
}

bool checkButtonPress() {
  if (digitalRead(button1) == LOW || digitalRead(button2) == LOW || digitalRead(button3) == LOW || digitalRead(button4) == LOW) {
    return true;
  }
  return false;
}

void resumeWithYellow() {
  switch (lastStep) {
    case 0:
      resumeYellowStep(greenLED1, yellowLED1, redLED1, redLED2, redLED3, redLED4);
      break;
    case 1:
      resumeYellowStep(greenLED2, yellowLED2, redLED2, redLED1, redLED3, redLED4);
      break;
    case 2:
      resumeYellowStep(greenLED3, yellowLED3, redLED3, redLED1, redLED2, redLED4);
      break;
    case 3:
      resumeYellowStep(greenLED4, yellowLED4, redLED4, redLED1, redLED2, redLED3);
      break;
  }
  resumeFromYellow = false;
}

void resumeYellowStep(int greenPin, int yellowPin, int redPin, int otherRed1, int otherRed2, int otherRed3) {
  turnAllOff();
  
  digitalWrite(otherRed1, HIGH);
  digitalWrite(otherRed2, HIGH);
  digitalWrite(otherRed3, HIGH);

  digitalWrite(yellowPin, HIGH);
  
  unsigned long startTime = millis();
  while (millis() - startTime < interval) {
    if (checkButtonPress()) return;
  }

  digitalWrite(yellowPin, LOW);
  digitalWrite(greenPin, HIGH);
}