let snakes = [];

let savedSnakes = [];

let population = 500;

let food;

let gridSize = 20;

let trainingSpeed = 1;

​

function setup() {

    createCanvas(400, 400);

    frameRate(10); // Adjust the frame rate as needed

    for(let i = 0; i < population; i++){

      append(snakes, new Snake());      

    }

}

​

function draw() {

    //background(0,0,0,5);

    background(100);

    for(let j = 0; j < trainingSpeed; j++){

    for(let i = snakes.length -1; i >= 0; i--){

      let snake = snakes[i];

      // Check for collisions

      if(snake.checkCollision()){

        savedSnakes.push(snakes.splice(i, 1)[0]);

      }

​

      // Update and display the snake

      snake.think();

      snake.update();

      snake.display();

      snake.eat();

    }

    if(snakes.length == 0){

      nextGeneration();

    }

    }

    // //Display score

    // fill(255);

    // textSize(16);

    // text(snake.body.length, 10, 20);

}

​

​

// function restart() {

//   for(let snake of snakes){

//     snake.body = [];

//     snake.body[0] = createVector(floor(width / 2), floor(height / 2));

//     snake.xDirection = 0;

//     snake.yDirection = -1;  

//   }

// }

​

​

function nextGeneration() {

  console.log('next generation');

  calculateFitness();

  for (let i = 0; i < population; i++) {

    snakes[i] = pickOne();

  }

  savedBirds = [];

}

​

function pickOne() {

  let index = 0;

  let r = random(1);

  while (r > 0) {

    r = r - savedSnakes[index].fitness;

    index++;

  }

  index--;

  let snake = savedSnakes[index];

  let child = new Snake(snake.brain);

  child.mutate();

  return child;

}

​

function calculateFitness() {

  let sum = 0;

  for (let snake of savedSnakes) {

    // sum += snake.score * (snake.body.length-1/2);

    sum += snake.score + (snake.body.length-1);

  }

  for (let snake of savedSnakes) {

    // snake.fitness = snake.score * (snake.body.length-1/2) / sum;

    snake.fitness = (snake.score + (snake.body.length-1)) / sum;

  }

}

​

​

​

​

​

​

// function keyPressed() {

//     //up

//     if (keyCode === 38) {

//       snake.turnUp();

//     }

//     //down

//     else if (keyCode === 40) {

//         snake.turnDown();

//     } 

//     //left

//     else if (keyCode === 37) {

//         snake.turnLeft();

//     } 

//     //right

//     else if (keyCode === 39) {

//         snake.turnRight(); 

//     }

// }

​

function createFood(snake) {

    let foodPos;

    let isFoodOnSnake = true;

​

    while (isFoodOnSnake) {

        const cols = floor(width / gridSize);

        const rows = floor(height / gridSize);

        const foodX = floor(random(cols)) * gridSize;

        const foodY = floor(random(rows)) * gridSize;

​

        // Check if the food is not on the snake's body

        isFoodOnSnake = false;

        for (let i = 0; i < snake.body.length; i++) {

            const segment = snake.body[i];

            if (foodX === segment.x && foodY === segment.y) {

                isFoodOnSnake = true;

                break;

            }

        }

​

        if (!isFoodOnSnake) {

            foodPos = createVector(foodX, foodY);

        }

    }

​

    return foodPos;

}

​

class Snake {

    constructor(brain) {

      this.body = [];

      this.body[0] = createVector(floor(width / 2), floor(height / 2));

      this.xDirection = 0;

      this.yDirection = 0;

      this.food = createFood(this);

      this.score = 0; 

      this.fitness = 0;

      //(input, hidden, output)

      //input head position x and y, snake velocity x and y, food location x and y. Need to add body

      //output up, down, left, right

      if(brain){

        this.brain = brain.copy();

      }else{

        this.brain = new NeuralNetwork(6, 12, 4);        

      }

    }

    think() {

    const head = this.body[this.body.length - 1].copy();

    const currentXDir = this.xDirection;

    const currentYDir = this.yDirection;

​

    // Calculate the inputs for the neural network

    let inputs = [];

    inputs[0] = head.x / width;

    inputs[1] = head.y / height;

    inputs[2] = currentXDir;

    inputs[3] = currentYDir;

    inputs[4] = this.food.x / width;

    inputs[5] = this.food.y / height;

​

    // Get the output from the neural network

    let output = this.brain.predict(inputs);

​

    // Determine the new direction based on the neural network's output

    let newDirectionX = currentXDir;

    let newDirectionY = currentYDir;

​

    if (output[0] > output[1] && output[0] > output[2] && output[0] > output[3]) {

        // Keep the current direction

      newDirectionX = (currentXDir === 0) ? -1 : currentXDir; // Prevent reversing

        newDirectionY = 0;

    } else if (output[1] > output[0] && output[1] > output[2] && output[1] > output[3]) {

        // Change direction to right

        newDirectionX = (currentXDir === 0) ? 1 : currentXDir; // Prevent reversing

        newDirectionY = 0;

    } else if (output[2] > output[0] && output[2] > output[1] && output[2] > output[3]) {

        // Change direction to up

        newDirectionX = 0;

        newDirectionY = (currentYDir === 0) ? -1 : currentYDir; // Prevent reversing

    } else if (output[3] > output[0] && output[3] > output[1] && output[3] > output[2]) {

        // Change direction to down

        newDirectionX = 0;

        newDirectionY = (currentYDir === 0) ? 1 : currentYDir; // Prevent reversing

    }

​

    // Set the new direction

    this.setDirection(newDirectionX, newDirectionY);

}

    mutate() {

      this.brain.mutate(0.1);

    }

    turnLeft(){

      if(this.xDirection != 1){

      //if(this.xDirection != 1 && this.yDireciton != 0){

        this.setDirection(-1, 0);

      //}

      }

    }

    turnRight(){

      if(this.xDirection != -1){

        this.setDirection(1, 0);                

      }

      //if(this.xDirection != -1 && this.yDireciton != 0){

      //}

    }  

    turnUp(){

      //if(this.xDirection != 0 && this.yDireciton != 1){

      if(this.yDirection != 1){

        this.setDirection(0, -1);

      }

      //}

    }

    turnDown(){

      //if(this.xDirection != 0 && this.yDireciton != -1){

      if(this.yDirection != -1){

        this.setDirection(0, 1); 

      }

      //}

    }

    setDirection(x, y) {

        this.xDirection = x;

        this.yDirection = y;

    }

​

    update() {

        this.score ++;  

        const head = this.body[this.body.length - 1].copy();

        this.body.shift(); // Remove the tail

​

        head.x += this.xDirection * gridSize;

        head.y += this.yDirection * gridSize;

​

        this.body.push(head); // Add a new head

    }

​

    grow() {

        const head = this.body[this.body.length - 1].copy();

        this.body.push(head);

    }

​

    eat() {

        const head = this.body[this.body.length - 1];

        if (head.x === this.food.x && head.y === this.food.y) {

            this.grow();

            this.food = createFood(this);

            return true;

        }

        return false;

    }

​

    checkCollision() {

      const head = this.body[this.body.length - 1];

      for (let i = 0; i < this.body.length - 2; i++) {

          const segment = this.body[i];

          if (head.x === segment.x && head.y === segment.y) {

              // Snake collided with itself

              return true;

          }

      }

      if (head.x >= width || head.x < 0 || head.y >= height || head.y < 0) {

          //console.log('Game Over - Wall Collision');

          //restart();

          return true;

      }

    }

    display() {

        fill(255, 10, 10);

        // Display the food  

        rect(this.food.x, this.food.y, gridSize, gridSize);

        // Color the head 

        //-todo: just make some eyes

        fill(0,255,255);

        noStroke(); 

        rect(this.body[this.body.length - 1].x, this.body[this.body.length - 1].y, gridSize, gridSize);

        // Color the body of the snake

        for (let i = 0; i < this.body.length -1; i++) {

            fill(0, 100, 0);

            noStroke();

          rect(this.body[i].x,this.body[i].y,gridSize,gridSize)

            //rect(this.body[i].x + 10*noise(this.body[i].x,this.body[i].y,frameCount), this.body[i].y + 10*noise(this.body[i].x,this.body[i].y,frameCount), gridSize, gridSize);

        }

    }

}