pathing_a_star.py

from __future__ import annotations
import pygame as pg
from pygame import Vector2 as vec2
import math
import heapq
import itertools


class Game:
  def __init__(self, width:int, height:int):
    pg.init()
    self.window = pg.display.set_mode((width, height))
    self.win_size = vec2(width, height)
    self.clock = pg.time.Clock()
    self.running = False
    self.dt = 0.0
    self.tile_size = 32

    self.tiles = [
      [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
      [0,1,1,1,1,1,1,1,1,1,1,1,1,1,1,0],
      [0,1,1,2,1,1,1,1,1,1,1,1,1,1,1,0],
      [0,1,1,3,1,1,1,1,1,1,1,1,1,1,1,0],
      [0,1,1,4,1,1,1,1,1,1,1,1,1,1,1,0],
      [0,1,1,5,1,1,1,1,1,1,1,1,1,1,1,0],
      [0,1,1,6,1,1,1,1,1,1,1,1,1,1,1,0],
      [0,1,1,7,1,1,1,1,1,1,1,1,1,1,1,0],
      [0,1,1,8,1,1,1,1,1,1,1,1,1,1,1,0],
      [0,1,1,9,1,1,1,1,1,1,1,1,1,1,1,0],
      [0,1,1,0,1,1,1,1,1,1,1,1,1,1,1,0],
      [0,1,1,0,1,1,1,1,1,1,1,1,1,1,1,0],
      [0,1,1,0,1,1,1,1,1,1,1,1,1,1,1,0],
      [0,1,1,1,1,1,1,1,1,1,1,1,1,1,1,0],
      [0,1,1,1,1,1,1,1,1,1,1,1,1,1,1,0],
      [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],
    ]
    self.player_pos = vec2(10, 10)
    self.movement_range = 6

    self.valid_tiles = get_valid_positions(
      self.player_pos,
      self.movement_range,
      self.tiles,
    )
    self.path = None


  def run(self):
    self.running = True
    while self.running:
      self.handle_events()
      self.update()
      self.draw()


  def handle_events(self):
    for event in pg.event.get():
      if (event.type == pg.QUIT or 
          (event.type == pg.KEYDOWN and
           event.key == pg.K_ESCAPE)):
        self.running = False
      elif event.type == pg.MOUSEBUTTONDOWN:
        pos = (
          event.pos[0] // self.tile_size,
          event.pos[1] // self.tile_size
        )
        if pos in self.valid_tiles: 
          self.player_pos = vec2(pos)
          self.valid_tiles = get_valid_positions(
            self.player_pos,
            self.movement_range,
            self.tiles,
          )
          self.path = None
      elif event.type == pg.MOUSEMOTION and self.valid_tiles:
        pos = (
          event.pos[0] // self.tile_size,
          event.pos[1] // self.tile_size
        )
        if pos in self.valid_tiles:
          self.path = a_star(
            tuple(map(int, self.player_pos)),
            pos,
            self.valid_tiles
          )
        else:
          self.path = None

  def update(self):
    self.dt = self.clock.tick(60) * 0.001
    fps = self.clock.get_fps()
    pg.display.set_caption(f'{fps:0.2f}')


  def draw(self):
    self.window.fill((30,20,20))

    for y, row in enumerate(self.tiles):
      for x, tile in enumerate(row):
        pg.draw.rect(
          self.window,
          ((255 - (tile * 28)) if tile else 0,) * 3,
          (
            (x * self.tile_size, y * self.tile_size),
            (self.tile_size, self.tile_size)
          )
        )

    surf = pg.Surface(self.win_size, pg.SRCALPHA)
    for pos in self.valid_tiles:
      pg.draw.rect(
        surf,
        (0,100,0,100),
        (
          (pos[0] * self.tile_size, pos[1] * self.tile_size),
          (self.tile_size, self.tile_size)
        )
      )

    pg.draw.rect(
      surf,
      (100,0,0),
      (self.player_pos.elementwise() * self.tile_size, (self.tile_size, self.tile_size))
    )

    if self.path:
      for a, b in itertools.pairwise(self.path):
        start = vec2(a).elementwise() * vec2(self.tile_size) + vec2(self.tile_size).elementwise() * 0.5
        end = vec2(b).elementwise() * vec2(self.tile_size) + vec2(self.tile_size).elementwise() * 0.5
        pg.draw.line(surf, (200,0,0), start, end)

    self.window.blit(surf, (0,0))
    pg.display.update()


def get_valid_positions(start, distance, board):
  valid_positions = []
  queue = [(list(map(int, start)), 0)]

  while queue:
    (x, y), dist = queue.pop(0)
    if (
      0 <= y < len(board) and
      0 <= x < len(board[0]) and
      board[y][x] != 0
    ):
      if (x, y) not in valid_positions:
        valid_positions.append((x, y))

      if dist < distance:
        for dy in range(-1, 2):
          for dx in range(-1, 2):
            if dx == 0 and dy == 0:
              continue
            nx, ny = x + dx, y + dy
            if (
              0 <= ny < len(board) and
              0 <= nx < len(board[0]) and
              board[ny][nx] != 0
            ):
              diagonal = dx != 0 and dy != 0
              new_dist = dist + board[ny][nx] + (0.5 if diagonal else 0.0)
              if new_dist <= distance:
                queue.append(((nx, ny), new_dist))
  return valid_positions


def heuristic(a, b):
  return math.sqrt((b[0] - a[0]) ** 2 + (b[1] - a[1]) ** 2)


def a_star(start, goal, graph):
  neighbors = [(0,1),(0,-1),(1,0),(-1,0),(1,1),(1,-1),(-1,1),(-1,-1)]

  open_list = []
  heapq.heappush(open_list, (0, start))

  came_from = {}
  cost_so_far = {}
  came_from[start] = None
  cost_so_far[start] = 0

  while open_list:
    current = heapq.heappop(open_list)[1]

    if current == goal:
      break

    for i, j in neighbors:
      next_node = (current[0] + i, current[1] + j)
      if next_node not in graph:
        continue

      new_cost = cost_so_far[current] + heuristic(current, next_node)

      if next_node not in cost_so_far or new_cost < cost_so_far[next_node]:
        cost_so_far[next_node] = new_cost
        priority = new_cost + heuristic(goal, next_node)
        heapq.heappush(open_list, (priority, next_node))
        came_from[next_node] = current

  if goal not in came_from:
    return None

  current = goal
  path = []
  while current != start:
    path.append(current)
    current = came_from[current]
  path.append(start)
  path.reverse()
  return path


Game(512, 512).run()