diff --git a/src/AI/Affinities.py b/src/AI/Affinities.py index d2e1ac6..b35d505 100644 --- a/src/AI/Affinities.py +++ b/src/AI/Affinities.py @@ -1,3 +1,8 @@ +import random + +from src.entities.Enums import Classifiers + + class Affinities: def __init__(self, food, water, rest): """ @@ -9,3 +14,15 @@ class Affinities: self.food = food self.water = water self.rest = rest + + def getWeigths(self): + return [self.food, self.water, self.rest] + + + +def pickWeightedAffinity(affinities: Affinities): + population = [Classifiers.FOOD, Classifiers.WATER, Classifiers.REST] + weights = affinities.getWeigths() + + return random.choices(population, weights) +