simulator.py

# simulator.py
import pygame
import imageio
import numpy as np
from robot import Robot

class Simulator:
    def __init__(self, arena_size=500):
        pygame.init()
        self.arena_size = arena_size
        self.window = pygame.display.set_mode((arena_size, arena_size))
        pygame.display.set_caption("Brownian Robot Simulation")
        self.clock = pygame.time.Clock()
        self.robot = Robot(arena_size)
        self.path = [self.robot.position.copy()]
        self.frames = []

    def run(self):
        running = True
        while running:
            self.clock.tick(60)  # 60 FPS

            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False

            self.robot.move()
            self.path.append(self.robot.position.copy())

            self._draw()

        pygame.quit()
        imageio.mimsave("brownian_motion.gif", self.frames, fps=30)

    def _draw(self):
        self.window.fill((255, 255, 255))  # white background

        # Draw the path
        if len(self.path) > 1:
            pygame.draw.lines(
                self.window, (200, 200, 200), False,
                [pos.astype(int) for pos in self.path], 1
            )

        # Draw the robot
        pygame.draw.circle(
            self.window, (0, 0, 255),
            self.robot.position.astype(int),
            self.robot.radius
        )

        pygame.display.flip()

        # Capture frame for GIF
        frame = pygame.surfarray.array3d(self.window)
        frame = np.transpose(frame, (1, 0, 2))  # Convert to (H, W, C)
        self.frames.append(frame)