139 lines
5.8 KiB
Python
139 lines
5.8 KiB
Python
import pygame
|
|
import numpy as np
|
|
from config import MazeConfig
|
|
from maze import render_static
|
|
|
|
class Visualizer:
|
|
def __init__(self, screen: pygame.Surface, grid: np.ndarray, cfg: MazeConfig):
|
|
self.screen = screen
|
|
self.grid = grid
|
|
self.cfg = cfg
|
|
self.maze_w = cfg.maze_cols * cfg.cell_size
|
|
self.maze_h = cfg.maze_rows * cfg.cell_size
|
|
self.hud_x = self.maze_w + 10
|
|
# Pre-render static maze
|
|
self.maze_surface = render_static(grid, cfg)
|
|
# Trail surface with per-pixel alpha
|
|
self.trail_surf = pygame.Surface((self.maze_w, self.maze_h), pygame.SRCALPHA)
|
|
self.trail_surf.fill((0, 0, 0, 0))
|
|
# Pre-allocate fade surface — reused every frame (avoids per-frame Surface creation)
|
|
self._fade_surf = pygame.Surface((self.maze_w, self.maze_h), pygame.SRCALPHA)
|
|
subtract = max(1, int(255 * (1.0 - cfg.trail_decay)))
|
|
self._fade_surf.fill((0, 0, 0, subtract))
|
|
# Pre-render static gradient bar for HUD
|
|
self._gradient_bar = self._make_gradient_bar(170)
|
|
# Font
|
|
pygame.font.init()
|
|
self.font = pygame.font.SysFont("monospace", cfg.hud_font_size)
|
|
self.font_sm = pygame.font.SysFont("monospace", cfg.hud_font_size - 4)
|
|
self.clock = pygame.time.Clock()
|
|
self._goal_flash = 0
|
|
self._first_goal = False
|
|
|
|
def _make_gradient_bar(self, width: int) -> pygame.Surface:
|
|
surf = pygame.Surface((width, 10))
|
|
for i in range(width):
|
|
t = i / width
|
|
surf.set_at((i, 0), (int(255 * (1 - t)), int(255 * t), 30))
|
|
for y in range(1, 10):
|
|
surf.set_at((i, y), (int(255 * (1 - t)), int(255 * t), 30))
|
|
return surf
|
|
|
|
def _rank_color(self, rank: int, total: int) -> tuple[int, int, int]:
|
|
t = rank / max(1, total - 1)
|
|
return (int(255 * t), int(255 * (1.0 - t)), 30)
|
|
|
|
def _fade_trails(self):
|
|
self.trail_surf.blit(self._fade_surf, (0, 0), special_flags=pygame.BLEND_RGBA_SUB)
|
|
|
|
def _draw_hud(self, gen: int, best_fitness: float, step: int, alive: int, reached: int):
|
|
# Dark sidebar background
|
|
sidebar_rect = pygame.Rect(self.maze_w, 0, 200, self.maze_h)
|
|
pygame.draw.rect(self.screen, (10, 10, 16), sidebar_rect)
|
|
|
|
lines = [
|
|
("GEN", f"{gen:>4}"),
|
|
("STEP", f"{step:>4}/{self.cfg.max_steps}"),
|
|
("ALIVE", f"{alive:>4}/{self.cfg.population}"),
|
|
("GOAL", f"{reached:>4}"),
|
|
("BEST", f"{best_fitness:>8.1f}"),
|
|
]
|
|
y = 20
|
|
accent = (0, 220, 200)
|
|
for label, value in lines:
|
|
lbl_surf = self.font_sm.render(label, True, (100, 100, 120))
|
|
val_surf = self.font.render(value, True, accent)
|
|
self.screen.blit(lbl_surf, (self.hud_x, y))
|
|
self.screen.blit(val_surf, (self.hud_x, y + 16))
|
|
y += 52
|
|
|
|
# Legend — pre-rendered gradient bar (no per-frame line drawing)
|
|
y += 20
|
|
legend_label = self.font_sm.render("RANK", True, (80, 80, 100))
|
|
self.screen.blit(legend_label, (self.hud_x, y))
|
|
y += 18
|
|
self.screen.blit(self._gradient_bar, (self.hud_x, y))
|
|
best_lbl = self.font_sm.render("best", True, (0, 255, 30))
|
|
wrst_lbl = self.font_sm.render("worst", True, (255, 0, 30))
|
|
self.screen.blit(best_lbl, (self.hud_x, y + 14))
|
|
self.screen.blit(wrst_lbl, (self.hud_x + 120, y + 14))
|
|
|
|
def render(self, agents, gen: int, step: int):
|
|
# Sort by fitness for ranking colors
|
|
sorted_agents = sorted(agents, key=lambda a: a.fitness)
|
|
|
|
# Check first goal
|
|
reached = sum(1 for a in agents if a.reached_goal)
|
|
if reached > 0 and not self._first_goal:
|
|
self._first_goal = True
|
|
self._goal_flash = 45
|
|
|
|
# Blit static maze
|
|
self.screen.blit(self.maze_surface, (0, 0))
|
|
|
|
# Fade and draw trails
|
|
self._fade_trails()
|
|
for rank, agent in enumerate(sorted_agents):
|
|
if agent.trail:
|
|
color = self._rank_color(rank, len(sorted_agents))
|
|
px, py = agent.trail[-1]
|
|
pygame.draw.circle(self.trail_surf, (*color, 130), (px, py), 3)
|
|
self.screen.blit(self.trail_surf, (0, 0))
|
|
|
|
# Draw agents on top
|
|
for rank, agent in enumerate(sorted_agents):
|
|
color = self._rank_color(rank, len(sorted_agents))
|
|
px = agent.col * self.cfg.cell_size + self.cfg.cell_size // 2
|
|
py = agent.row * self.cfg.cell_size + self.cfg.cell_size // 2
|
|
radius = 5 if agent.reached_goal else 3
|
|
pygame.draw.circle(self.screen, color, (px, py), radius)
|
|
if agent.reached_goal:
|
|
pygame.draw.circle(self.screen, (255, 255, 255), (px, py), radius + 2, 1)
|
|
|
|
# Goal flash overlay
|
|
if self._goal_flash > 0:
|
|
gx = (self.cfg.maze_cols - 1) * self.cfg.cell_size + self.cfg.cell_size // 2
|
|
gy = (self.cfg.maze_rows - 1) * self.cfg.cell_size + self.cfg.cell_size // 2
|
|
alpha = int(200 * self._goal_flash / 45)
|
|
flash_surf = pygame.Surface((self.cfg.cell_size, self.cfg.cell_size), pygame.SRCALPHA)
|
|
flash_surf.fill((255, 220, 0, alpha))
|
|
self.screen.blit(flash_surf, (gx - self.cfg.cell_size//2, gy - self.cfg.cell_size//2))
|
|
self._goal_flash -= 1
|
|
|
|
# HUD
|
|
best_fit = max((a.fitness for a in agents), default=0.0)
|
|
alive = sum(1 for a in agents if not a.reached_goal)
|
|
self._draw_hud(gen, best_fit, step, alive, reached)
|
|
|
|
pygame.display.flip()
|
|
|
|
def reset_trails(self):
|
|
self.trail_surf.fill((0, 0, 0, 0))
|
|
self._first_goal = False
|
|
|
|
def new_maze(self, grid, cfg):
|
|
self.grid = grid
|
|
self.cfg = cfg
|
|
self.maze_surface = render_static(grid, cfg)
|
|
self.reset_trails()
|