main.py

import matplotlib

matplotlib.use("Agg")

import matplotlib.backends.backend_agg as agg

import pylab

from math import floor

import pygame
pygame.init()
from random import randint, random


WINDOW_X, WINDOW_Y = 1920, 1080
CENTER_X, CENTER_Y = WINDOW_X // 2, WINDOW_Y // 2
CENTER_RADIUS = 400

FPS = 999
SPEED = 1

GRAY_COLOR = (50, 50, 50)

WINDOW = pygame.display.set_mode((WINDOW_X, WINDOW_Y), pygame.FULLSCREEN) # pygame.FULLSCREEN
FONT = pygame.font.Font(None, 32)

CLOCK = pygame.time.Clock()

ENV_X = 0.0
ENV_Y = 0.0
ENV_SPEED_CHANGE = 0.1

BORN_MODULO = 300
MUTATION_RATE = 1
AGENTS_LIMIT = 10000
AGENTS_COUNT = 1000
STAGE = 0

agents = []
babies = []

class Agent:
    def __init__(self, x, y, ax, ay, mutation_chance):
        self.x = x
        self.y = y
        self.ax = ax
        self.ay = ay
        self.mut_chance = mutation_chance


    def step(self):
        self.x += self.ax
        self.y += self.ay

        self.x += ENV_X
        self.y += ENV_Y

    def born(self):
        babies.append(
            Agent(
                self.x,
                self.y,
                self.ax + self.mut_chance * randint(-1, 1) * random(),
                self.ay + self.mut_chance * randint(-1, 1) * random(),
                self.mut_chance + self.mut_chance *  randint(-1, 1) * random(),
            )
        )

    def is_alive(self):
        # Flip the y-axis
        flipped_y = WINDOW_Y - self.y

        # Check if the agent is within the circle
        if (self.x - CENTER_X) ** 2 + (flipped_y - CENTER_Y) ** 2 < CENTER_RADIUS ** 2:
            return True
        return False


def start(AGENTS_COUNT):
    for i in range(AGENTS_COUNT):
        agents.append(
            Agent(
                randint(0, WINDOW_X),
                randint(0, WINDOW_Y),
                randint(0, 2),
                randint(0, 2),
                random()
            )
        )

def agents_step():
    mid_mut_rate = 0
    condition = STAGE % BORN_MODULO == 0
    for agent in agents:
        agent.step()

        mid_mut_rate += agent.mut_chance

        if not agent.is_alive():
            agents.remove(agent)

        if condition and len(agents) + len(babies) < AGENTS_LIMIT:
            agent.born()

    if len(agents) != 0:
        mid_mut_rate /= len(agents)
    else:
        mid_mut_rate = 0
    graph_mutation_rates.append(mid_mut_rate)
    agents.extend(babies)
    babies.clear()


def check_events():
    global SPEED, ENV_X, ENV_Y
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            exit()
        if event.type == pygame.KEYDOWN:
            key = event.key
            if key == pygame.K_q:
                exit()
            if key == pygame.K_1:
                SPEED = 1
            if key == pygame.K_2:
                SPEED = 5
            if key == pygame.K_3:
                SPEED = 25
            if key == pygame.K_4:
                SPEED = 125
            if key == pygame.K_d:
                ENV_X += ENV_SPEED_CHANGE
            if key == pygame.K_a:
                ENV_X -= ENV_SPEED_CHANGE
            if key == pygame.K_w:
                ENV_Y += ENV_SPEED_CHANGE
            if key == pygame.K_s:
                ENV_Y -= ENV_SPEED_CHANGE


    ENV_X = round(ENV_X, 6)
    ENV_Y = round(ENV_Y, 6)


screen = pygame.display.get_surface()
surf = None

graph_agents_count = []
graph_mutation_rates = []



def draw_objects():
    WINDOW.fill((200, 200, 200))


    def draw_demographic_graph():

        def calculate_demographic_graph():
            global surf
            if STAGE % 200 != 0 and STAGE > 1:
                screen.blit(surf, (0, 150))
                return
            fig = pylab.figure(
                figsize=[4, 4],  # Inches
                dpi=120,  # 100 dots per inch, so the resulting buffer is 400x400 pixels
            )
            ax = fig.gca()
            ax.plot(graph_agents_count[::-1][:1000][::-1])

            canvas = agg.FigureCanvasAgg(fig)
            canvas.draw()
            renderer = canvas.get_renderer()
            raw_data = renderer.tostring_rgb()

            size = canvas.get_width_height()

            surf = pygame.image.fromstring(raw_data, size, "RGB")
            screen.blit(surf, (0, 150))

        calculate_demographic_graph()

    def draw_mutation_rate_graph():

        def calculate_mutation_rate_graph():
            global surf
            if STAGE % 200 != 0 and STAGE > 1:
                screen.blit(surf, (0, 600))
                return
            fig = pylab.figure(
                figsize=[4, 4],  # Inches
                dpi=120,  # 100 dots per inch, so the resulting buffer is 400x400 pixels
            )
            ax = fig.gca()
            ax.plot(graph_mutation_rates[::-1][:10000][::-1])

            canvas = agg.FigureCanvasAgg(fig)
            canvas.draw()
            renderer = canvas.get_renderer()
            raw_data = renderer.tostring_rgb()

            size = canvas.get_width_height()

            surf = pygame.image.fromstring(raw_data, size, "RGB")
            screen.blit(surf, (0, 600))

        calculate_mutation_rate_graph()


    def draw_big_circle():
        pygame.draw.circle(
            WINDOW,
            (50, 50, 50),
            (CENTER_X, CENTER_Y),
            CENTER_RADIUS
        )

    def draw_agents():
        for agent in agents:
            pygame.draw.circle(
                WINDOW,
                (0, 0, 0),
                (agent.x, agent.y),
                3
            )

    # draw_demographic_graph()
    draw_mutation_rate_graph()
    draw_big_circle()
    draw_agents()

def draw_text():
    def draw_current_stage():
        text = "CYCLE: " + str(floor(STAGE / 100))
        WINDOW.blit(
            FONT.render(
                text,
                True,
                (255, 0, 0)
            ),
            (20, 20),
        )
    def draw_agents_count():
        text = "hAgents: " + str(len(agents) / 100)
        WINDOW.blit(
            FONT.render(
                text,
                True,
                (255, 0, 0)
            ),
            (20, 80),
        )
    def draw_current_env():
        text = "MOV_X: " + str(ENV_X)
        WINDOW.blit(
            FONT.render(
                text,
                True,
                (255, 0, 0)
            ),
            (20, 40),
        )
        text = "MOV_Y: " + str(ENV_Y)
        WINDOW.blit(
            FONT.render(
                text,
                True,
                (255, 0, 0)
            ),
            (20, 60),
        )

    draw_current_stage()
    draw_current_env() # test
    draw_agents_count()

def clear():
    pygame.display.flip()

start(AGENTS_COUNT)
while True:
    ENV_X += 0.0001
    ENV_X = round(ENV_X, 6)

    STAGE += 1
    if STAGE % SPEED == 0:
        agents_step()
    check_events()
    draw_objects()
    draw_text()
    clear()
    CLOCK.tick(FPS)
    if STAGE % 100 == 0:
        graph_agents_count.append(len(agents))
        text = \
        f"""
        STAGE: {STAGE}   AGENTS:{len(agents)}
        """
        print(text)