population.py

import copy
from settings import Settings
from snake import Snake
from food import Food
from inputs import Inputs
from NeuralNetwork import NeuralNetwork
import numpy as np
from p5 import *
import random
import heapq
import bisect
import itertools
from cprofiler import profile 


class Population:
    def __init__(self, settings, onGenerationComplete=None):
        self.settings: Settings = settings
        self.onGenerationComplete = onGenerationComplete
        self.FoodArr: list[Food] = []
        self.SnakeArr: list[Snake] = []
        self.InputArr: list[Inputs] = []
        self.NNArr: list[NeuralNetwork] = []
        self.snakesAlive = settings.numberOfSnakes
        self.genBestIndex = 0
    def setup(self, seedGenomes=None):
        seedGenomes = seedGenomes or []
        for i in range(self.settings.numberOfSnakes):
            self.FoodArr.append(Food(self.settings))
            self.SnakeArr.append(Snake(self.settings, self.FoodArr[i]))
            self.InputArr.append(Inputs(self.settings, self.SnakeArr[i]))
            self.NNArr.append(NeuralNetwork(
                self.settings,
                self.InputArr[i],
                initializeWeights=not seedGenomes,
            ))
        if seedGenomes:
            self.seedPopulation(seedGenomes)

    def seedPopulation(self, seedGenomes):
        if len(seedGenomes) == 1:
            founder = seedGenomes[0]
        elif len(seedGenomes) == 2:
            parent1 = self.NNArr[0].copyWeightsFrom(seedGenomes[0])
            parent2 = self.NNArr[1].copyWeightsFrom(seedGenomes[1])
            founder = NeuralNetwork(
                self.settings,
                self.InputArr[0],
                initializeWeights=False,
            ).crossover(parent1, parent2)
        else:
            raise ValueError("Seed a population with one or two brain genomes.")

        for i, brain in enumerate(self.NNArr):
            brain.copyWeightsFrom(founder)
            if i > 0:
                brain.mutate()

    def run(self):
        for i in range(self.settings.numberOfSnakes):
            self.SnakeArr[i].update()
            if(self.SnakeArr[i].alive == True):
                self.FoodArr[i].update()
                self.InputArr[i].generateInputs()
                self.NNArr[i].controlSnake(self.SnakeArr[i])
            self.snakesAlive = self.settings.numberOfSnakesAlive
            if(self.snakesAlive == 0):
                self.calcBestSnakeBrainIndex()
                print("gen bests ------ index: {} eats: {} score: {}".format(self.genBestIndex, self.SnakeArr[self.genBestIndex].total, self.SnakeArr[self.genBestIndex].fitness))
                print("champion -------- index: {} seed: {} score: {}".format(self.settings.globalBestIndex, self.settings.currentFoodSeed, self.settings.globalBestScore))
                if self.onGenerationComplete is not None:
                    self.onGenerationComplete(
                        self.NNArr[self.genBestIndex],
                        self.SnakeArr[self.genBestIndex],
                        self.settings,
                    )
                self.startNextGeneration()
                return

        self.showTopSnakes()

    def getVisibleSnakeIndexes(self):
        aliveIndexes = [
            i for i, snake in enumerate(self.SnakeArr)
            if snake.alive
        ]
        visibleLimit = self.settings.visibleSnakeLimit
        if len(aliveIndexes) <= visibleLimit:
            return aliveIndexes

        return heapq.nlargest(
            visibleLimit,
            aliveIndexes,
            key=lambda i: (self.SnakeArr[i].fitness, -i),
        )

    def showTopSnakes(self):
        for i in self.getVisibleSnakeIndexes():
            self.FoodArr[i].show()
            self.SnakeArr[i].show()

    def startNextGeneration(self):
        print("starting next gen")
        #increment generation
        self.settings.generation += 1

        #do algorithm for genetics
        self.createNewBrains()

        #reset scores
        self.settings.reset()
        self.resetSnakesAndFood()

    def createNewBrains(self):
        #create a New NN from Old NN that has been mutated
        self.calcBestSnakeBrainIndex()
        parentBrains = self.NNArr
        childBrains = []
        self.cumulativeFitness = list(itertools.accumulate(
            snake.fitness for snake in self.SnakeArr
        ))
        for i in range(self.settings.numberOfSnakes):
            # Keep this scenario's winner. It is reevaluated with every other
            # brain on the next generation's newly seeded food scenario.
            if(i != self.genBestIndex):
                #create babies and match them with snakes based on score
                childBrains.append(self.naturalSelection(i, parentBrains))
            else:
                childBrains.append(parentBrains[i])
        self.NNArr = childBrains

    def resetSnakesAndFood(self):
        for i in range(self.settings.numberOfSnakes):
            # self.FoodArr[i].spawnAtRandomLocation()
            self.SnakeArr[i].resetSnake()

    def naturalSelection(self, index, parentBrains=None):
        if parentBrains is None:
            parentBrains = self.NNArr
        #create a new NeuralNetwork
        childNN = NeuralNetwork(
            self.settings,
            self.InputArr[index],
            initializeWeights=False,
        )
        
        # select two brains from the existing pool to mate
        selectedIndex1 = self.selectNNFromSnakeFitness()
        selectedIndex2 = self.selectNNFromSnakeFitness()
        childNN = childNN.crossover(parentBrains[selectedIndex1], parentBrains[selectedIndex2])
        childNN = childNN.mutate()
        return(childNN)
                
    def selectNNFromSnakeFitness(self):
        # has test in test_1.py
        if self.settings.totalFitness <= 0:
            return random.randrange(self.settings.numberOfSnakes)
        rand = random.randint(1, self.settings.totalFitness)
        cumulativeFitness = getattr(self, "cumulativeFitness", None)
        if cumulativeFitness is None:
            cumulativeFitness = list(itertools.accumulate(
                snake.fitness for snake in self.SnakeArr
            ))
        return bisect.bisect_left(cumulativeFitness, rand)

    def calcBestSnakeBrainIndex(self):
        genBestScore = -1
        for i in range(self.settings.numberOfSnakes):
            if(self.SnakeArr[i].fitness > genBestScore):
                genBestScore = self.SnakeArr[i].fitness
                self.genBestIndex = i
        self.settings.globalBestScore = genBestScore
        self.settings.globalBestIndex = self.genBestIndex