COSC2968|COSC3053 Foundations of Artificial Intelligence for STEM: Assignment 2

astar.py

import pygame
import itertools    
import math
import time
import random

# Constants
MAZE_GENERATION_VISUALIZATION = True  # Set to True to visualize the maze generation process
MAZE_GAP = 2  # The gap between the walls in the maze
# 
GOAL_COUNT = 3  # Number of goals in the maze aside from the end
MOVEMENT_TYPE = "4"  # 4 or 8 directions
# Visualization settings
DELAY = 0 # Delay between each step in seconds
WIDTH = 21 # Width of the maze
HEIGHT = 21 # Height of the maze
CELL_SIZE = 20 # Size of each cell in pixels

class Node:
    def __init__(self, parent=None, position=None):
        self.parent = parent
        self.position = position

        self.g = 0
        self.h = 0
        self.f = 0

    def __eq__(self, other):
        return self.position == other.position

def draw_maze(maze, current_node, closed_list, open_list, start, end):
    # Define colors
    BLACK = (0, 0, 0)
    WHITE = (255, 255, 255)
    GREY = (200, 200, 200)
    BLUE = (0, 255, 255)
    RED = (255, 0, 0)
    GREEN = (0, 255, 0)

    # Draw the maze
    for i in range(HEIGHT):
        for j in range(WIDTH):
            color = WHITE 
            if maze[i][j] == 1:
                color = BLACK
            elif (i, j) == current_node.position:
                color = BLUE
            elif (i, j) == start:
                color = RED
            elif (i, j) == end:
                color = GREEN
            elif maze[i][j] == 2: # The additional goals
                color = (255, 255, 0)
                pygame.draw.rect(screen, color, pygame.Rect(j * CELL_SIZE, i * CELL_SIZE, CELL_SIZE, CELL_SIZE))
                # Additionally write a number displaying what goal number it is
                font = pygame.font.Font(None, 24)
                text_surface = font.render(str(goals.index((i, j))), True, BLACK)
                screen.blit(text_surface, (j * CELL_SIZE + 5, i * CELL_SIZE + 5))
                continue
            elif (i, j) in [node.position for node in closed_list]:
                dist = math.sqrt((i - end[0]) ** 2 + (j - end[1]) ** 2)
                highest_dist = highest_dist = math.sqrt(WIDTH ** 2 + HEIGHT ** 2)
                # The color should be red and green the closer it is to the goal
                color = (255 * (dist / highest_dist), 255 * (1 - dist / highest_dist), 0)
                # Draw the rectangle with a border
                pygame.draw.rect(screen, WHITE, pygame.Rect(j * CELL_SIZE, i * CELL_SIZE, CELL_SIZE, CELL_SIZE))
                pygame.draw.rect(screen, color, pygame.Rect(j * CELL_SIZE + 2.5, i * CELL_SIZE + 2.5, CELL_SIZE - 5, CELL_SIZE - 5))
                continue
            elif (i, j) in [node.position for node in open_list]:
                color = GREY

            pygame.draw.rect(screen, color, pygame.Rect(j * CELL_SIZE, i * CELL_SIZE, CELL_SIZE, CELL_SIZE))

    pygame.display.flip()
    
def draw_infobox(start, end, open_list, closed_list, distance_matrix, total_distance=0):
    infobox_x = WIDTH * CELL_SIZE
    infobox_width = 250
    infobox_height = HEIGHT * CELL_SIZE
    infobox_rect = pygame.Rect(infobox_x, 0, infobox_width, infobox_height)
    
    # Draw infobox background
    pygame.draw.rect(screen, (0, 0, 0), infobox_rect)
    pygame.draw.rect(screen, (255, 255, 255), infobox_rect, 2)  # Add border for clarity
    
    # Font settings
    font = pygame.font.Font(None, 24)
    text_color = (255, 255, 255)
    
    # Draw text information
    info = [
        f"Start: {start}",
        f"Goal: {end}",
        f"Open List Size: {len(open_list)}",
        f"Closed List Size: {len(closed_list)}"
    ]
    
    y_offset = 10
    for line in info:
        text_surface = font.render(line, True, text_color)
        screen.blit(text_surface, (infobox_x + 10, y_offset))
        y_offset += 30

    # Visualize Distance Matrix
    screen.blit(font.render("Distance Matrix", True, text_color), (infobox_x + 10, y_offset + 10))
    
    matrix_x = infobox_x + 10
    matrix_y = y_offset + 30
    matrix_cell_size = 30
    
    for i in range(len(distance_matrix)):
        for j in range(len(distance_matrix[i])):
            color = (0, 255, 0)
            if distance_matrix[i][j] == float('inf'):
                color = (255, 0, 0)
            pygame.draw.rect(screen, color, pygame.Rect(matrix_x + j * matrix_cell_size, matrix_y + i * matrix_cell_size, matrix_cell_size, matrix_cell_size))
            text_surface = font.render("{:.0f}".format(distance_matrix[i][j] / 2), True, text_color)
            screen.blit(text_surface, (matrix_x + j * matrix_cell_size + 5, matrix_y + i * matrix_cell_size + 5))
    
    path_distance_text = f"Final Path Distance: {total_distance:.2f}"
    text_surface = font.render(path_distance_text, True, text_color)
    screen.blit(text_surface, (infobox_x + 10, matrix_y + len(distance_matrix) * matrix_cell_size + 20))
    
    pygame.display.flip()
    
def draw_triangle(screen, start_pos, direction, color=(0, 0, 0)):
    TRIANGLE_SIZE = CELL_SIZE / 2
    
    # Normalize the direction vector
    dir_length = math.sqrt(direction[0]**2 + direction[1]**2)
    direction = (direction[0] / dir_length, direction[1] / dir_length)
    
    # Calculate the angle of the direction
    angle = math.atan2(direction[1], direction[0])
    
    # Calculate the position of the end of the triangle's base
    base_pos = (
        start_pos[0] - direction[0] * (TRIANGLE_SIZE * 0.5),
        start_pos[1] - direction[1] * (TRIANGLE_SIZE * 0.5)
    )
    
    # Calculate the vertices of the triangle
    left_vertex = (
        base_pos[0] + TRIANGLE_SIZE * 0.5 * math.cos(angle + math.pi / 2),
        base_pos[1] + TRIANGLE_SIZE * 0.5 * math.sin(angle + math.pi / 2)
    )
    right_vertex = (
        base_pos[0] + TRIANGLE_SIZE * 0.5 * math.cos(angle - math.pi / 2),
        base_pos[1] + TRIANGLE_SIZE * 0.5 * math.sin(angle - math.pi / 2)
    )
    tip_vertex = (
        start_pos[0] + direction[0] * (TRIANGLE_SIZE * 0.5),
        start_pos[1] + direction[1] * (TRIANGLE_SIZE * 0.5)
    )
    
    # Draw the filled triangle
    pygame.draw.polygon(screen, color, [tip_vertex, left_vertex, right_vertex])
    # Draw the tip of the triangle
    pygame.draw.circle(screen, color, (int(tip_vertex[0]), int(tip_vertex[1])), 2)

# Drawing the optimal path
# Drawing the optimal path
def draw_path(path, total_distance=None):
    # Wipe the screen, keeping the maze
    screen.fill((0, 0, 0))
    
    # Draw the maze
    for i in range(HEIGHT):
        for j in range(WIDTH):
            color = (255, 255, 255) 
            if maze[i][j] == 1:
                color = (0, 0, 0)
            elif maze[i][j] == 2: # The additional goals
                color = (255, 255, 0)
            pygame.draw.rect(screen, color, pygame.Rect(j * CELL_SIZE, i * CELL_SIZE, CELL_SIZE, CELL_SIZE))
    
    draw_infobox(start, path[len(path) - 1], [], [], distance_matrix, total_distance)
    
    # Initialize a dictionary to track the number of visits to each cell
    visit_count = {}
    
    for idx in range(1, len(path)):
        position = path[idx]
        prev_position = path[idx - 1]
        
        # Track visits
        if position not in visit_count:
            visit_count[position] = 0
        visit_count[position] += 1
        
        # Modify the color based on the number of visits
        visit_color = (0, 255 // visit_count[position], 255)  # Adjust the green component
        
        # Draw the path segment
        direction = (position[0] - prev_position[0], position[1] - prev_position[1])
        draw_triangle(screen, (prev_position[1] * CELL_SIZE + CELL_SIZE // 2, prev_position[0] * CELL_SIZE + CELL_SIZE // 2), direction, (255, 0, 0))
        
        # Do not draw on top of goals
        if maze[position[0]][position[1]] == 2:
            continue
        
        pygame.draw.rect(screen, visit_color, pygame.Rect(position[1] * CELL_SIZE, position[0] * CELL_SIZE, CELL_SIZE, CELL_SIZE))
        
        pygame.display.flip()
        time.sleep(DELAY + 0.05)
            
# Draw a random maze with a given size.
def generate_maze():
    if (HEIGHT - 1) % MAZE_GAP != 0 or (WIDTH - 1) % MAZE_GAP != 0:
        raise ValueError("HEIGHT and WIDTH must be 1 greater than a multiple of MAZE_GAP.")
    
    # Fill the maze with walls
    maze = [[1 for _ in range(WIDTH)] for _ in range(HEIGHT)]
    start = (0, 0)
    
    stack = [start]
    visited = set()
    came_from = {}  # Dictionary to keep track of the path
    
    while stack:
        current = stack[-1]
        visited.add(current)
        
        # Clear the current cell and the surrounding path cells
        for i in range(MAZE_GAP):
            for j in range(MAZE_GAP):
                if current[0] + i < HEIGHT and current[1] + j < WIDTH:
                    maze[current[0] + i][current[1] + j] = 0
        
        unvisited_neighbors = []
        for i, j in [(0, MAZE_GAP + 1), (0, -(MAZE_GAP + 1)), (MAZE_GAP + 1, 0), (-(MAZE_GAP + 1), 0)]:
            neighbor = (current[0] + i, current[1] + j)
            if 0 <= neighbor[0] < HEIGHT and 0 <= neighbor[1] < WIDTH and neighbor not in visited:
                unvisited_neighbors.append(neighbor)
        
        if unvisited_neighbors:
            next_cell = random.choice(unvisited_neighbors)
            stack.append(next_cell)
            
            # Clear the path between current and next cell
            if next_cell[0] == current[0]:  # Horizontal move
                step = 1 if next_cell[1] > current[1] else -1
                for col in range(current[1] + step, next_cell[1] + step, step):
                    for i in range(MAZE_GAP):
                        if current[0] + i < HEIGHT and col < WIDTH:
                            maze[current[0] + i][col] = 0
            else:  # Vertical move
                step = 1 if next_cell[0] > current[0] else -1
                for row in range(current[0] + step, next_cell[0] + step, step):
                    for j in range(MAZE_GAP):
                        if row < HEIGHT and current[1] + j < WIDTH:
                            maze[row][current[1] + j] = 0
            
            came_from[next_cell] = current  # Track the path
        else:
            stack.pop()
        
        # Draw the generation process if visualization is enabled
        if MAZE_GENERATION_VISUALIZATION and draw_maze and Node:
            draw_maze(maze, Node(None, current), [], [], start, (HEIGHT - 1, WIDTH - 1))
    
    # Randomly place the goals in the maze
    goals = [start]
    maze[0][0] = 2
    for _ in range(GOAL_COUNT):
        goal = (random.randint(0, HEIGHT - 1), random.randint(0, WIDTH - 1))
        while maze[goal[0]][goal[1]] == 1:
            goal = (random.randint(0, HEIGHT - 1), random.randint(0, WIDTH - 1))
        maze[goal[0]][goal[1]] = 2
        goals.append(goal)
    
    print(maze)
    print(goals)
    return maze, start, goals

# A* search algorithm for a grid map
# The goal is to find the shortest path from the start to the end
def astar(maze, start, end):
    # Create start and end node
    start_node = Node(None, start)
    start_node.g = start_node.h = start_node.f = 0
    end_node = Node(None, end)
    end_node.g = end_node.h = end_node.f = 0

    # Initialize both open and closed list
    open_list = []
    closed_list = []

    # Add the start node
    open_list.append(start_node)

    # Loop until you find the end
    while open_list:
        time.sleep(DELAY)  # Sleep for 0.05 seconds

        # Get the current node
        current_node = open_list[0]
        current_index = 0
        for index, item in enumerate(open_list):
            if item.f < current_node.f:
                current_node = item
                current_index = index

        # Pop current off open list, add to closed list
        open_list.pop(current_index)
        closed_list.append(current_node)

        # Check if we found the goal
        if current_node == end_node:
            path = []
            current = current_node
            while current is not None:
                path.append(current.position)
                current = current.parent
            return path[::-1]  # Return reversed path

        # Swap between two maze types
        # If the maze is a grid map, we can move in 4 or 8 directions
        # If the maze is a graph, we can move to any other node.
        # Determine what direction can we navigate
        if MOVEMENT_TYPE == "4":
            directions = [(0, -1), (0, 1), (-1, 0), (1, 0)]
        elif MOVEMENT_TYPE == "8":
            directions = [(0, -1), (0, 1), (-1, 0), (1, 0), (-1, -1), (-1, 1), (1, -1), (1, 1)]

        # Expansion: Generate children
        children = []
        for new_position in directions:
            # Get node position
            node_position = (current_node.position[0] + new_position[0], current_node.position[1] + new_position[1])

            # Make sure within range
            if node_position[0] > (len(maze) - 1) or node_position[0] < 0 or node_position[1] > (len(maze[0]) - 1) or node_position[1] < 0:
                continue

            # Make sure walkable terrain
            if maze[node_position[0]][node_position[1]] != 0 and maze[node_position[0]][node_position[1]] != 2:  # 0 is walkable
                continue

            # Create new node
            new_node = Node(current_node, node_position)
            children.append(new_node)
                    
        # Loop through children
        for child in children:
            if child in closed_list:
                continue

            # Calculate the g: movement cost to move from the starting point to a given node
            # As we are using a grid map, the cost is always 1
            child.g = current_node.g + 1
            # Calculate the h: heuristic cost. Here we use the Chebyshev distance as the heuristic function
            child.h = max(abs(child.position[0] - end_node.position[0]), abs(child.position[1] - end_node.position[1]))     
            # Calculate the f: f = g + h, the total cost
            child.f = child.g + child.h

            # Check whether the child is in the open list
            if child in open_list:
                index = open_list.index(child)
                # If the child's f
                if open_list[index].f > child.f:
                    open_list[index] = child
            else:
                open_list.append(child)

        # Draw the maze
        draw_maze(maze, current_node, closed_list, open_list, start, end)
        draw_infobox(start, end, open_list, closed_list, distance_matrix)
    return None

def calculate_route_distance(permutation, distance_matrix):
    # Calculate the total distance of the route based on the distance matrix.
    total_distance = 0
    for i in range(len(permutation) - 1):
        total_distance += distance_matrix[permutation[i]][permutation[i+1]]
    # Add distance to return to the start
    total_distance += distance_matrix[permutation[-1]][permutation[0]]
    return total_distance

def find_shortest_path(distance_matrix):
    # Find the shortest path visiting all goals and returning to the start.
    num_goals = len(distance_matrix)
    goals_indices = list(range(num_goals))
    min_distance = float('inf')
    best_permutation = None
    
    for permutation in itertools.permutations(goals_indices):
        current_distance = calculate_route_distance(permutation, distance_matrix)
        if current_distance < min_distance:
            min_distance = current_distance
            best_permutation = permutation
    
    return best_permutation, min_distance / 2 # Divide by 2 to get the distance of the path

if __name__ == '__main__':
    # Set up display
    pygame.init()
    pygame.display.set_caption('A* Pathfinding Visualization')
    window_size = (WIDTH * CELL_SIZE + 250, HEIGHT * CELL_SIZE)

    screen = pygame.display.set_mode(window_size)
    # White background
    screen.fill((255, 255, 255))

    # Generate a random maze
    maze, start, goals = generate_maze()
    #maze = [[2, 0, 1, 0, 0, 0, 0], [0, 0, 1, 0, 0, 0, 0], [0, 0, 1, 1, 1, 1, 0], [0, 0, 0, 0, 0, 0, 0], [2, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 0], [0, 2, 0, 0, 0, 0, 0]]
    #start = (0, 0)
    #goals = [(0, 0), (6, 1), (6, 1), (4, 0)]
    # Print goals
    print("Goals:", goals)

    # Calculate the distance of the most optimized path for every goal pair
    distance_matrix = [[float('inf')] * (GOAL_COUNT + 1) for _ in range(GOAL_COUNT + 1)]
    paths = [[None] * (GOAL_COUNT + 1) for _ in range(GOAL_COUNT + 1)]
    
    checkedGoalPair = []
    for i, goal in enumerate(goals):
        for j, goal2 in enumerate(goals):
            # If the goals are the same, the distance is 0
            if goal == goal2:
                distance_matrix[i][j] = 0
            elif (i, j) not in checkedGoalPair and (j, i) not in checkedGoalPair:
                print(f"Checking distance between {goal} and {goal2}")
                path = astar(maze, goal, goal2)
                if path:
                    distance = len(path)
                    # Store the distance symmetrically
                    distance_matrix[i][j] = distance
                    distance_matrix[j][i] = distance
                     # Store the path symmetrically
                    paths[i][j] = path
                    # Reverse the path to get the path from goal2 to goal
                    paths[j][i] = path[::-1]
                else:
                    distance_matrix[i][j] = float('inf')
                    distance_matrix[j][i] = float('inf')
                # Add this pair to the checked pairs
                checkedGoalPair.append((i, j))
                    
    print("Distance Matrix:" , distance_matrix)
    # Find the most optimal path using the brute force algorithm
    best_route, min_distance = find_shortest_path(distance_matrix)
    print("Best route:", best_route)
    print("Minimum distance:", min_distance)
    print("Paths:", paths)

    # Construct the path based on the best route using a* algorithm
    full_path = []
    

    # Start with the first segment of the path
    full_path.extend(paths[best_route[0]][best_route[1]])

    # Construct the path based on the best route using a* algorithm
    full_path = []
    visited_positions = set()

    # Start with the first segment of the path
    initial_segment = paths[best_route[0]][best_route[1]]
    full_path.extend(initial_segment)
    visited_positions.update(initial_segment)

    # Iterate through the best_route and append each segment, ensuring no backtracking
    for i in range(1, len(best_route) - 1):
        segment = paths[best_route[i]][best_route[i + 1]]
        if segment:
            # Append only the non-overlapping portion of the segment
            non_overlapping_segment = [pos for pos in segment if pos not in visited_positions]
            full_path.extend(non_overlapping_segment)
            visited_positions.update(non_overlapping_segment)

    # Draw the constructed path
    draw_path(full_path, min_distance)
          
    # Wait for the user to close the window
    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False