#include <LedControl.h> 
#include <Servo.h>

#define numberOfDisplays 4

#define JOYSTICK_X A1
#define JOYSTICK_Y A0
#define JOYSTICK_BTN 2

#define SERVO_PIN 9
#define SERVO_INIT 0
#define SERVO_MID 1
#define SERVO_MIN 2
#define SERVO_MAX 3

LedControl lc1 = LedControl(12, 13, 11, numberOfDisplays); 
LedControl lc2 = LedControl(10, 7, 8, numberOfDisplays); 
LedControl lc3 = LedControl(6, 4, 5, numberOfDisplays); 
LedControl lc4 = LedControl(3, A3, A2, numberOfDisplays); 

Servo myservo;

const int snakeMaxLength = 1024;
struct SnakeSegment {
  int row;
  int col;
};
struct SnakeSegment snake[snakeMaxLength]; 
int snakeLength = 3;
bool gameRunning = false;  
bool gameOver = false;  
enum Direction { 
  UP, 
  DOWN, 
  LEFT, 
  RIGHT 
};
Direction snakeDirection = RIGHT; 

int beanrow, beancol;

int joystickXValue = 0;
int joystickYValue = 0;

unsigned long lastServoMoveTime = 0;
int servoState = SERVO_INIT; 

void setup() {
  Serial.begin(9600);
  for (int i = 0; i < numberOfDisplays; i++) {
    lc1.shutdown(i, false);  
    lc2.shutdown(i, false);  
    lc3.shutdown(i, false);  
    lc4.shutdown(i, false);  
    lc1.setIntensity(i, 8);  
    lc2.setIntensity(i, 8); 
    lc3.setIntensity(i, 8);  
    lc4.setIntensity(i, 8); 
    lc1.clearDisplay(i);  
    lc2.clearDisplay(i);    
    lc3.clearDisplay(i);  
    lc4.clearDisplay(i);     
  }
  pinMode(JOYSTICK_X, INPUT);
  pinMode(JOYSTICK_Y, INPUT);
  pinMode(JOYSTICK_BTN, INPUT_PULLUP);
  myservo.attach(SERVO_PIN);
  myservo.write(90);
  startGame();
}

void loop() {
  if (gameRunning == true && gameOver == false) {
    moveSnake();  
    checkCollision();
    servoshake();
    delay(300);
  } 
  else if (gameOver == true) {
    for (int i = 0; i < numberOfDisplays; i++) {
      lc1.clearDisplay(i);  
      lc2.clearDisplay(i);    
      lc3.clearDisplay(i);  
      lc4.clearDisplay(i);     
    }
  }
  if(digitalRead(JOYSTICK_BTN) == LOW){
    startGame();
  }
}

void startGame() {
  snakeLength = 3;
  snake[0] = {15, 15};
  snake[1] = {15, 14};
  snake[2] = {15, 13};
  snakeDirection = RIGHT;
  generateBean();
  gameRunning = true;
  gameOver = false;
}

void generateBean() {
  randomSeed(millis());
  beanrow = random(0, 32);
  beancol = random(0, 32);
}

void moveSnake() {
  SnakeSegment tmp[snakeMaxLength];
  for(int j = 0; j < snakeLength; j++){
    tmp[j].row = snake[j].row;
    tmp[j].col = snake[j].col;
  }
  int original_snakeDirection = snakeDirection;
  readJoystick();
  if(snakeDirection == RIGHT && original_snakeDirection != LEFT){
    for(int i = snakeLength - 1; i > 0; i--){
      snake[i].row = tmp[i-1].row;
      snake[i].col = tmp[i-1].col;
    }
    snake[0].col++;
    updateDisplay();
  }
  else if(snakeDirection == LEFT && original_snakeDirection != RIGHT){
    for(int i = snakeLength - 1; i > 0; i--){
      snake[i].row = tmp[i-1].row;
      snake[i].col = tmp[i-1].col;
    }
    snake[0].col--;
    updateDisplay();
  }
  else if(snakeDirection == UP && original_snakeDirection != DOWN){
    for(int i = snakeLength - 1; i > 0; i--){
      snake[i].row = tmp[i-1].row;
      snake[i].col = tmp[i-1].col;
    }
    snake[0].row--;
    updateDisplay();
  }
  else if(snakeDirection == DOWN && original_snakeDirection != UP){
    for(int i = snakeLength - 1; i > 0; i--){
      snake[i].row = tmp[i-1].row;
      snake[i].col = tmp[i-1].col;
    }
    snake[0].row++;
    updateDisplay();
  }

  if (snake[0].row == beanrow && snake[0].col == beancol) {
    snakeLength++;
    generateBean();
  }
}

void checkCollision(){
  if (snake[0].row < 0 || snake[0].row>= 32 || snake[0].col < 0 || snake[0].col >= 32) {
    gameOver = true;
  }
  for (int i = 1; i < snakeLength; i++) {
    if (snake[0].row == snake[i].row && snake[0].col == snake[i].col) {
      gameOver = true;
      break;
    }
  }
}

void servoshake() {
  if (snakeLength >= 10) {
    unsigned long currentTime = millis();
    if (currentTime - lastServoMoveTime > 1000) { 
      lastServoMoveTime = currentTime;
      switch (servoState) {
        case SERVO_INIT:
          myservo.write(90);
          servoState = SERVO_MID;
          break;
        case SERVO_MID:
          myservo.write(0);
          servoState = SERVO_MIN;
          break;
        case SERVO_MIN:
          myservo.write(180);
          servoState = SERVO_MAX;
          break;
        case SERVO_MAX:
          myservo.write(90);
          servoState = SERVO_MID;
          break;
      }
    }
  }
}

void updateDisplay() {
  for (int i = 0; i < numberOfDisplays; i++) {
    lc1.clearDisplay(i);  
    lc2.clearDisplay(i);    
    lc3.clearDisplay(i);  
    lc4.clearDisplay(i);     
  }
  for (int i = 0; i < snakeLength; i++) {
    setLed(snake[i].row, snake[i].col, true);
  }
  setLed(beanrow, beancol, true);
}

void setLed(int row, int col, bool state) {
  int displayIndex1, displayIndex2, localrow, localcol;
  if (row < 8) {
    displayIndex1 = 0;
    if (col < 8) {
      displayIndex2 = 3;
      localcol = col;
      localrow = row;
    } else if (col < 16) {
      displayIndex2 = 2;
      localcol = col - 8;
      localrow = row;
    } else if (col < 24) {
      displayIndex2 = 1;
      localcol = col - 16;
      localrow = row;
    } else {
      displayIndex2 = 0;
      localcol = col - 24;
      localrow = row;
    }
  } 
  else if (row < 16) {
    displayIndex1 = 1;
    if (col < 8) {
      displayIndex2 = 3;
      localcol = col;
      localrow = row - 8;
    } else if (col < 16) {
      displayIndex2 = 2;
      localcol = col - 8;
      localrow = row - 8;
    } else if (col < 24) {
      displayIndex2 = 1;
      localcol = col - 16;
      localrow = row - 8;
    } else {
      displayIndex2 = 0;
      localcol = col - 24;
      localrow = row - 8;
    }
  } 
  else if (row < 24) {
    displayIndex1 = 2;
    if (col < 8) {
      displayIndex2 = 3;
      localcol = col;
      localrow = row - 16;
    } else if (col < 16) {
      displayIndex2 = 2;
      localcol = col - 8;
      localrow = row - 16;
    } else if (col < 24) {
      displayIndex2 = 1;
      localcol = col - 16;
      localrow = row - 16;
    } else {
      displayIndex2 = 0;
      localcol = col - 24;
      localrow = row - 16;
    }
  } 
  else{
    displayIndex1 = 3;
    if (col < 8) {
      displayIndex2 = 3;
      localcol = col;
      localrow = row - 24;
    } else if (col < 16) {
      displayIndex2 = 2;
      localcol = col - 8;
      localrow = row - 24;
    } else if (col < 24) {
      displayIndex2 = 1;
      localcol = col - 16;
      localrow = row - 24;
    } else {
      displayIndex2 = 0;
      localcol = col - 24;
      localrow = row - 24;
    }
  } 
  switch(displayIndex1) {
    case 0:
      lc1.setLed(displayIndex2, localrow, localcol, state);
      break;
    case 1:
      lc2.setLed(displayIndex2, localrow, localcol, state);
      break;
    case 2:
      lc3.setLed(displayIndex2, localrow, localcol, state);
      break;
    case 3:
      lc4.setLed(displayIndex2, localrow, localcol, state);
      break;
  }
}

void readJoystick() {
  int joystickXValue = analogRead(JOYSTICK_X);
  int joystickYValue = analogRead(JOYSTICK_Y);
  if (joystickXValue < 300 && snakeDirection != LEFT) {
    snakeDirection = RIGHT;
  } 
  else if (joystickXValue > 700 && snakeDirection != RIGHT) {
    snakeDirection = LEFT;
  } 
  else if (joystickYValue < 300 && snakeDirection != UP) {
    snakeDirection = DOWN;
  } 
  else if (joystickYValue > 700 && snakeDirection != DOWN) {
    snakeDirection = UP;
  }
}