astar and BMSSP Lite

main.py

"""
Grid-based demonstration of BMSSP Lite vs Classic Dijkstra.
English only. No Cyrillic.
"""

import pyray as pr
import random
import math
import heapq
from dataclasses import dataclass, field

# --- CONSTANTS & COLORS ---
COLS = 40
ROWS = 40
CELL_SIZE = 12

COLOR_NORMAL = pr.RAYWHITE
COLOR_SWAMP = pr.Color(210, 180, 140, 255)  # Tan/Brown
COLOR_WALL = pr.DARKGRAY
COLOR_START = pr.BLUE
COLOR_END = pr.RED
COLOR_FRONTIER = pr.YELLOW
COLOR_SETTLED = pr.GREEN
COLOR_PATH = pr.MAGENTA
COLOR_CHEAP = pr.SKYBLUE

# --- DATA STRUCTURES ---


@dataclass
class RunStats:
    heap_push: int = 0
    heap_pop: int = 0
    cheap_relax: int = 0


class AlgoState:
    def __init__(self, cols, rows):
        self.dist = {}
        self.settled = set()
        self.in_heap = set()
        self.parent = {}
        self.path = []
        self.stats = RunStats()
        self.phase = 'Waiting...'
        self.done = False
        self.active_cells = set()

        for x in range(cols):
            for y in range(rows):
                self.dist[(x, y)] = float('inf')


# --- GRID GENERATION ---


def generate_grid(cols, rows):
    grid = {}
    for x in range(cols):
        for y in range(rows):
            r = random.random()
            if r < 0.20:
                grid[(x, y)] = float('inf')  # Wall
            elif r < 0.40:
                grid[(x, y)] = 5.0  # Swamp
            else:
                grid[(x, y)] = 1.0  # Normal
    return grid


def get_neighbors(u, grid, cols, rows):
    x, y = u
    neighbors = []
    directions = [(0, -1), (0, 1), (-1, 0), (1, 0)]
    for dx, dy in directions:
        nx, ny = x + dx, y + dy
        if 0 <= nx < cols and 0 <= ny < rows:
            weight = grid[(nx, ny)]
            if weight != float('inf'):
                neighbors.append(((nx, ny), weight))
    return neighbors


# --- ALGORITHMS (GENERATORS FOR ANIMATION) ---


def reconstruct_path(state, start, end):
    curr = end
    while curr in state.parent:
        state.path.append(curr)
        curr = state.parent[curr]
        if curr == start:
            state.path.append(start)
            break


def dijkstra_classic_gen(grid, start, end, state):
    state.dist[start] = 0
    heap = [(0, start)]
    state.stats.heap_push += 1
    state.in_heap.add(start)
    state.phase = 'Searching (Heap)'

    while heap:
        d, u = heapq.heappop(heap)
        state.stats.heap_pop += 1
        state.in_heap.discard(u)

        if u in state.settled:
            continue

        state.settled.add(u)
        state.active_cells = {u}

        if u == end:
            break

        for v, w in get_neighbors(u, grid, COLS, ROWS):
            nd = d + w
            if nd < state.dist[v]:
                state.dist[v] = nd
                state.parent[v] = u
                heapq.heappush(heap, (nd, v))
                state.stats.heap_push += 1
                state.in_heap.add(v)
        yield

    reconstruct_path(state, start, end)
    state.active_cells = set()
    state.phase = 'Done!'
    state.done = True
    yield


def bmssp_lite_gen(grid, start, end, state):
    n = COLS * ROWS
    INF = float('inf')
    state.dist[start] = 0

    logn = max(2.0, math.log2(max(n, 2)))
    k = max(1, round(logn ** (1 / 3)))
    t = max(1, round(logn ** (2 / 3)))

    in_heap_best = {}
    global_heap = []

    def push_global(v, d):
        if d < in_heap_best.get(v, INF):
            in_heap_best[v] = d
            heapq.heappush(global_heap, (d, v))
            state.stats.heap_push += 1
            state.in_heap.add(v)

    push_global(start, 0)

    while global_heap:
        d, u = heapq.heappop(global_heap)
        state.stats.heap_pop += 1
        state.in_heap.discard(u)

        if u in state.settled or d > state.dist[u]:
            continue

        if u == end:
            state.settled.add(u)
            break

        # 1. Cheap Relaxations (Bellman-Ford style)
        state.phase = 'Cheap Relax (No Heap)'
        W = {u}
        frontier = {u}
        touched1 = []

        for _ in range(k):
            nxt = set()
            state.active_cells = set(frontier)
            yield

            for curr_u in frontier:
                for v, w in get_neighbors(curr_u, grid, COLS, ROWS):
                    state.stats.cheap_relax += 1
                    nd = state.dist[curr_u] + w
                    if nd < state.dist[v]:
                        state.dist[v] = nd
                        state.parent[v] = curr_u
                        touched1.append(v)
                        nxt.add(v)
            if not nxt:
                break
            W |= nxt
            if len(W) > k * 4:
                break
            frontier = nxt

        for v in touched1:
            if v not in state.settled:
                push_global(v, state.dist[v])

        # 2. Local Heap (Base Case)
        state.phase = 'Local Heap'
        local_heap = [(state.dist[u], u)]
        state.stats.heap_push += 1
        local_settled = []
        seen_local = set()
        touched2 = []

        while local_heap and len(local_settled) < t + 1:
            ld, lu = heapq.heappop(local_heap)
            state.stats.heap_pop += 1

            if lu in seen_local or ld > state.dist[lu]:
                continue

            seen_local.add(lu)
            local_settled.append(lu)
            state.active_cells = {lu}
            yield

            if lu == end:
                break

            for v, w in get_neighbors(lu, grid, COLS, ROWS):
                nd = ld + w
                if nd < state.dist[v]:
                    state.dist[v] = nd
                    state.parent[v] = lu
                    heapq.heappush(local_heap, (nd, v))
                    state.stats.heap_push += 1
                    touched2.append(v)

        for v in touched2:
            if v not in state.settled:
                push_global(v, state.dist[v])

        for s in local_settled:
            state.settled.add(s)

        state.settled.add(u)

        if end in state.settled:
            break

    reconstruct_path(state, start, end)
    state.active_cells = set()
    state.phase = 'Done!'
    state.done = True
    yield


# --- DRAWING ---


def draw_grid_map(offset_x, offset_y, grid, state, start_pos, end_pos, title):
    pr.draw_text(title, offset_x, offset_y - 60, 20, pr.DARKGRAY)
    pr.draw_text(f'Phase: {state.phase}', offset_x, offset_y - 35, 18, pr.MAROON)

    heap_ops = state.stats.heap_push + state.stats.heap_pop
    pr.draw_text(f'Heap Ops (Expensive): {heap_ops}', offset_x, offset_y + ROWS * CELL_SIZE + 10, 18, pr.RED)
    pr.draw_text(
        f'Cheap Relax (O(1)): {state.stats.cheap_relax}', offset_x, offset_y + ROWS * CELL_SIZE + 30, 18, pr.BLUE
    )

    for x in range(COLS):
        for y in range(ROWS):
            px = offset_x + x * CELL_SIZE
            py = offset_y + y * CELL_SIZE

            weight = grid[(x, y)]
            color = COLOR_NORMAL

            if weight == float('inf'):
                color = COLOR_WALL
            elif weight == 5.0:
                color = COLOR_SWAMP

            if (x, y) in state.settled:
                color = COLOR_SETTLED
            elif (x, y) in state.in_heap:
                color = COLOR_FRONTIER

            if (x, y) in state.active_cells:
                color = COLOR_CHEAP if 'Cheap' in state.phase else pr.MAGENTA

            if (x, y) in state.path:
                color = COLOR_PATH
            if (x, y) == start_pos:
                color = COLOR_START
            if (x, y) == end_pos:
                color = COLOR_END

            pr.draw_rectangle(px, py, CELL_SIZE - 1, CELL_SIZE - 1, color)


# --- MAIN LOOP ---


def main():
    pr.init_window(1100, 650, 'Dijkstra vs BMSSP Lite - Grid Pathfinding')
    pr.set_target_fps(60)

    grid = generate_grid(COLS, ROWS)

    start_pos = (2, 2)
    end_pos = (COLS - 3, ROWS - 3)

    # Ensure start and end are not walls
    grid[start_pos] = 1.0
    grid[end_pos] = 1.0

    state_classic = AlgoState(COLS, ROWS)
    state_bmssp = AlgoState(COLS, ROWS)

    gen_classic = None
    gen_bmssp = None

    steps_per_frame = 10

    dragging_start = False
    dragging_end = False

    left_offset_x = 40
    right_offset_x = 560
    offset_y = 80

    while not pr.window_should_close():
        mouse_pos = pr.get_mouse_position()

        # --- INPUT HANDLING ---
        if gen_classic is None:  # Only allow dragging if not running
            # Check which grid the mouse is over to calculate grid coordinates
            grid_x = -1
            grid_y = -1

            if (
                left_offset_x <= mouse_pos.x < left_offset_x + COLS * CELL_SIZE
                and offset_y <= mouse_pos.y < offset_y + ROWS * CELL_SIZE
            ):
                grid_x = int((mouse_pos.x - left_offset_x) // CELL_SIZE)
                grid_y = int((mouse_pos.y - offset_y) // CELL_SIZE)
            elif (
                right_offset_x <= mouse_pos.x < right_offset_x + COLS * CELL_SIZE
                and offset_y <= mouse_pos.y < offset_y + ROWS * CELL_SIZE
            ):
                grid_x = int((mouse_pos.x - right_offset_x) // CELL_SIZE)
                grid_y = int((mouse_pos.y - offset_y) // CELL_SIZE)

            is_valid_cell = 0 <= grid_x < COLS and 0 <= grid_y < ROWS and grid.get((grid_x, grid_y)) != float('inf')

            if pr.is_mouse_button_pressed(pr.MOUSE_BUTTON_LEFT):
                if (grid_x, grid_y) == start_pos:
                    dragging_start = True
                elif (grid_x, grid_y) == end_pos:
                    dragging_end = True

            if pr.is_mouse_button_released(pr.MOUSE_BUTTON_LEFT):
                dragging_start = False
                dragging_end = False

            if dragging_start and is_valid_cell and (grid_x, grid_y) != end_pos:
                start_pos = (grid_x, grid_y)
            if dragging_end and is_valid_cell and (grid_x, grid_y) != start_pos:
                end_pos = (grid_x, grid_y)

            if pr.is_key_pressed(pr.KEY_SPACE):
                state_classic = AlgoState(COLS, ROWS)
                state_bmssp = AlgoState(COLS, ROWS)
                gen_classic = dijkstra_classic_gen(grid, start_pos, end_pos, state_classic)
                gen_bmssp = bmssp_lite_gen(grid, start_pos, end_pos, state_bmssp)

        if pr.is_key_pressed(pr.KEY_ENTER):
            gen_classic = None
            gen_bmssp = None
            state_classic = AlgoState(COLS, ROWS)
            state_bmssp = AlgoState(COLS, ROWS)

        # --- ANIMATION STEPS ---
        if gen_classic and not state_classic.done:
            for _ in range(steps_per_frame):
                try:
                    next(gen_classic)
                except StopIteration:
                    break

        if gen_bmssp and not state_bmssp.done:
            for _ in range(steps_per_frame):
                try:
                    next(gen_bmssp)
                except StopIteration:
                    break

        # --- DRAWING ---
        pr.begin_drawing()
        pr.clear_background(pr.RAYWHITE)

        # Draw UI Instructions
        instructions = 'Drag Start (Blue) / End (Red) | Press SPACE to Start | Press ENTER to Reset'
        pr.draw_text(instructions, 1100 // 2 - pr.measure_text(instructions, 20) // 2, 10, 20, pr.BLACK)

        # Draw Grids
        draw_grid_map(left_offset_x, offset_y, grid, state_classic, start_pos, end_pos, 'Classic Dijkstra O(m log n)')
        draw_grid_map(right_offset_x, offset_y, grid, state_bmssp, start_pos, end_pos, 'BMSSP Lite O(m log^(2/3) n)')

        # Draw Legend
        legend_y = offset_y + ROWS * CELL_SIZE + 60
        pr.draw_rectangle(left_offset_x, legend_y, 15, 15, COLOR_NORMAL)
        pr.draw_text('Normal (Cost 1)', left_offset_x + 20, legend_y, 15, pr.DARKGRAY)

        pr.draw_rectangle(left_offset_x + 150, legend_y, 15, 15, COLOR_SWAMP)
        pr.draw_text('Swamp (Cost 5)', left_offset_x + 170, legend_y, 15, pr.DARKGRAY)

        pr.draw_rectangle(left_offset_x + 300, legend_y, 15, 15, COLOR_WALL)
        pr.draw_text('Wall (Block)', left_offset_x + 320, legend_y, 15, pr.DARKGRAY)

        pr.end_drawing()

    pr.close_window()


if __name__ == '__main__':
    main()