main.py

import sys
import pygame
from constants import (
    WINDOW_WIDTH, WINDOW_HEIGHT, GRID_AREA_WIDTH, GRID_SIZE, CELL_SIZE,
    COLOR_BG, FPS, STEPS_PER_FRAME, AppState,
    NUM_DYNAMIC_OBSTACLES, DYNAMIC_OBS_RADIUS, DYNAMIC_OBS_SPEED,
)
from grid import Grid
from state import StateMachine
from renderer import Renderer
from ui import Button, AlgorithmToggle, AnalyticsPanel
from obstacles import create_dynamic_obstacles
from algorithms.bfs import bfs
from algorithms.dijkstra import dijkstra
from algorithms.astar import astar
from algorithms.greedy import greedy
from algorithms.rrt import rrt
from algorithms.rrt_star import rrt_star


def is_rrt_type(algo_name):
    return algo_name in ("RRT", "RRT*")


def create_algorithm_generator(algo_name, grid, dynamic_obstacles):
    """Create and return the appropriate algorithm generator."""
    start = grid.start
    end = grid.end
    if start is None or end is None:
        return None

    if algo_name == "BFS":
        return bfs(grid, start, end, dynamic_obstacles)
    elif algo_name == "Dijkstra":
        return dijkstra(grid, start, end, dynamic_obstacles)
    elif algo_name == "A*":
        return astar(grid, start, end, dynamic_obstacles)
    elif algo_name == "Greedy":
        return greedy(grid, start, end, dynamic_obstacles)
    elif algo_name in ("RRT", "RRT*"):
        start_px = grid.cell_to_pixel_center(*start)
        end_px = grid.cell_to_pixel_center(*end)
        bounds = (0, 0, GRID_AREA_WIDTH, WINDOW_HEIGHT)
        if algo_name == "RRT":
            return rrt(grid, start_px, end_px, bounds, dynamic_obstacles)
        else:
            return rrt_star(grid, start_px, end_px, bounds, dynamic_obstacles)
    return None


def compute_path_length(path, algo_name, cell_size):
    """Compute the length of a path."""
    if not path:
        return 0
    if is_rrt_type(algo_name):
        # Continuous path: sum of Euclidean distances
        import math
        length = 0
        for i in range(len(path) - 1):
            dx = path[i + 1][0] - path[i][0]
            dy = path[i + 1][1] - path[i][1]
            length += math.hypot(dx, dy)
        return round(length / cell_size, 1)  # Normalize to grid units
    else:
        return len(path) - 1  # Grid steps


def main():
    pygame.init()
    screen = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
    pygame.display.set_caption("Pathfinding Visualizer: BFS / Dijkstra / A* / Greedy / RRT / RRT*")
    clock = pygame.time.Clock()

    # Core objects
    grid = Grid(GRID_SIZE, GRID_SIZE, CELL_SIZE)
    sm = StateMachine()
    renderer = Renderer(screen, grid, CELL_SIZE)

    # Dynamic obstacles
    obs_bounds = (0, 0, GRID_AREA_WIDTH, WINDOW_HEIGHT)
    dynamic_obs = create_dynamic_obstacles(
        NUM_DYNAMIC_OBSTACLES, obs_bounds, DYNAMIC_OBS_RADIUS, DYNAMIC_OBS_SPEED)

    # Algorithm state
    algo_gen = None
    last_step = None
    path_blocked = False

    # UI elements
    panel = AnalyticsPanel(GRID_AREA_WIDTH, 0, 400, WINDOW_HEIGHT)
    toggle = AlgorithmToggle(820, 16, 360, 30)

    def on_start():
        nonlocal algo_gen, last_step, path_blocked
        if grid.start is not None and grid.end is not None:
            algo_gen = create_algorithm_generator(toggle.selected, grid, dynamic_obs)
            last_step = None
            path_blocked = False
            sm.transition(AppState.RUNNING)
            toggle.enabled = False

    def on_clear():
        nonlocal algo_gen, last_step, path_blocked
        grid.clear()
        sm.reset()
        algo_gen = None
        last_step = None
        path_blocked = False
        toggle.enabled = True
        panel.clear_results()

    btn_start = Button(820, 90, 170, 36, "Start", on_start)
    btn_clear = Button(1010, 90, 170, 36, "Clear", on_clear)

    # Determine steps per frame based on algorithm
    def get_steps_per_frame():
        if is_rrt_type(toggle.selected):
            return STEPS_PER_FRAME * 4  # RRT/RRT* need more iterations
        return STEPS_PER_FRAME

    fps_font = pygame.font.SysFont("consolas", 12)

    running = True
    mouse_held = False

    while running:
        # === EVENT HANDLING ===
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
                continue

            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    running = False
                    continue

            btn_start.handle_event(event)
            btn_clear.handle_event(event)
            toggle.handle_event(event)
            panel.handle_event(event)

            # Grid click handling
            if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
                mx, my = event.pos
                if mx < GRID_AREA_WIDTH and my < WINDOW_HEIGHT:
                    row, col = grid.pixel_to_cell(mx, my)
                    sm.handle_grid_click(row, col, grid)
                    if sm.click_count > 2:
                        mouse_held = True

            if event.type == pygame.MOUSEBUTTONUP and event.button == 1:
                mouse_held = False

        # Drag to draw obstacles
        if mouse_held and sm.state == AppState.SETUP and sm.click_count > 2:
            mx, my = pygame.mouse.get_pos()
            if mx < GRID_AREA_WIDTH and my < WINDOW_HEIGHT:
                row, col = grid.pixel_to_cell(mx, my)
                if grid.in_bounds(row, col):
                    if (row, col) != grid.start and (row, col) != grid.end:
                        grid.set_obstacle(row, col)

        # === UPDATE ===
        # Step the algorithm
        if sm.state == AppState.RUNNING and algo_gen is not None:
            spf = get_steps_per_frame()
            for _ in range(spf):
                try:
                    last_step = next(algo_gen)
                except StopIteration:
                    break
                if last_step.done:
                    sm.transition(AppState.DONE)
                    toggle.enabled = True
                    # Record results
                    path_len = 0
                    if last_step.path:
                        path_len = compute_path_length(
                            last_step.path, toggle.selected, CELL_SIZE)
                    panel.update_result(
                        toggle.selected,
                        last_step.elapsed_ms,
                        last_step.nodes_visited,
                        path_len,
                    )
                    break

        # Move dynamic obstacles always so user can see them
        for obs in dynamic_obs:
            obs.update()

        # Detect if a dynamic obstacle is blocking the found path
        if sm.state == AppState.DONE and last_step and last_step.path:
            path_blocked = False
            is_rrt = is_rrt_type(toggle.selected)
            for obs in dynamic_obs:
                if is_rrt:
                    if obs.intersects_path_continuous(last_step.path):
                        path_blocked = True
                        break
                else:
                    if obs.intersects_path_grid(last_step.path, CELL_SIZE):
                        path_blocked = True
                        break

        # Update button states — allow Start when we have endpoints (SETUP or DONE)
        has_endpoints = grid.start is not None and grid.end is not None
        btn_start.enabled = has_endpoints and sm.state in (AppState.SETUP, AppState.DONE)

        # === RENDER ===
        screen.fill(COLOR_BG)

        # Draw grid area
        renderer.draw_all(last_step, toggle.selected, dynamic_obs, path_blocked)

        # Draw panel
        status = sm.state.name
        if path_blocked:
            status = "PATH BLOCKED!"
        panel.draw(screen,
                   state_text=status,
                   algo_text=toggle.selected)

        # Draw UI
        toggle.draw(screen)
        btn_start.draw(screen)
        btn_clear.draw(screen)

        # FPS counter
        fps = clock.get_fps()
        fps_surf = fps_font.render(f"FPS: {fps:.0f}", True, (100, 100, 100))
        screen.blit(fps_surf, (GRID_AREA_WIDTH + 10, WINDOW_HEIGHT - 20))

        pygame.display.flip()
        clock.tick(FPS)

    pygame.quit()
    sys.exit()


if __name__ == "__main__":
    main()