import pygame
import random
from queue import PriorityQueue
import time

WIDTH = 600
ROWS = 20
CELL_SIZE = WIDTH // ROWS
FPS = 60
DELAY = 0

WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
BLUE = (0, 0, 255)
GRAY = (200, 200, 200)
YELLOW = (255, 255, 0)

class Cell:
    def __init__(self, row, col):
        self.row = row
        self.col = col
        self.x = col * CELL_SIZE
        self.y = row * CELL_SIZE
        self.walls = [True, True, True, True]  # top, right, bottom, left
        self.visited = False

    def draw(self, win):
        if self.visited:
            pygame.draw.rect(win, WHITE, (self.x, self.y, CELL_SIZE, CELL_SIZE))
        if self.walls[0]: pygame.draw.line(win, BLACK, (self.x, self.y), (self.x+CELL_SIZE, self.y), 2)
        if self.walls[1]: pygame.draw.line(win, BLACK, (self.x+CELL_SIZE, self.y), (self.x+CELL_SIZE, self.y+CELL_SIZE), 2)
        if self.walls[2]: pygame.draw.line(win, BLACK, (self.x+CELL_SIZE, self.y+CELL_SIZE), (self.x, self.y+CELL_SIZE), 2)
        if self.walls[3]: pygame.draw.line(win, BLACK, (self.x, self.y+CELL_SIZE), (self.x, self.y), 2)

    def highlight(self, win, color):
        pygame.draw.rect(win, color, (self.x, self.y, CELL_SIZE, CELL_SIZE))


def generate_maze(grid, current):
    stack = []
    current.visited = True

    while True:
        neighbors = []
        r, c = current.row, current.col
        if r > 0 and not grid[r-1][c].visited: neighbors.append(grid[r-1][c])
        if r < ROWS-1 and not grid[r+1][c].visited: neighbors.append(grid[r+1][c])
        if c > 0 and not grid[r][c-1].visited: neighbors.append(grid[r][c-1])
        if c < ROWS-1 and not grid[r][c+1].visited: neighbors.append(grid[r][c+1])

        if neighbors:
            next_cell = random.choice(neighbors)
            stack.append(current)
            remove_walls(current, next_cell)
            next_cell.visited = True
            current = next_cell
        elif stack:
            current = stack.pop()
        else:
            break

def remove_walls(a, b):
    dx = a.col - b.col
    dy = a.row - b.row
    if dx == 1:
        a.walls[3] = False
        b.walls[1] = False
    elif dx == -1:
        a.walls[1] = False
        b.walls[3] = False
    if dy == 1:
        a.walls[0] = False
        b.walls[2] = False
    elif dy == -1:
        a.walls[2] = False
        b.walls[0] = False

def heuristic(a, b):
    return abs(a.row - b.row) + abs(a.col - b.col)

def astar(grid, start, end, win):
    count = 0
    open_set = PriorityQueue()
    open_set.put((0, count, start))
    came_from = {}
    g_score = {cell: float("inf") for row in grid for cell in row}
    g_score[start] = 0
    f_score = {cell: float("inf") for row in grid for cell in row}
    f_score[start] = heuristic(start, end)
    open_set_hash = {start}

    while not open_set.empty():
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                return []

        current = open_set.get()[2]
        open_set_hash.remove(current)

        if current == end:
            path = []
            while current in came_from:
                current = came_from[current]
                path.append(current)
            return path

        for neighbor in get_neighbors(grid, current):
            temp_g = g_score[current] + 1
            if temp_g < g_score[neighbor]:
                came_from[neighbor] = current
                g_score[neighbor] = temp_g
                f_score[neighbor] = temp_g + heuristic(neighbor, end)
                if neighbor not in open_set_hash:
                    count += 1
                    open_set.put((f_score[neighbor], count, neighbor))
                    open_set_hash.add(neighbor)

        # visualize
        draw_grid(win, grid)
        for cell in open_set_hash:
            cell.highlight(win, GREEN)
        pygame.display.update()
        time.sleep(DELAY)

    return []

def get_neighbors(grid, cell):
    neighbors = []
    r, c = cell.row, cell.col
    if not cell.walls[0] and r > 0: neighbors.append(grid[r-1][c])
    if not cell.walls[1] and c < ROWS-1: neighbors.append(grid[r][c+1])
    if not cell.walls[2] and r < ROWS-1: neighbors.append(grid[r+1][c])
    if not cell.walls[3] and c > 0: neighbors.append(grid[r][c-1])
    return neighbors

# --- DRAWING ---
def draw_grid(win, grid):
    win.fill(WHITE)
    for row in grid:
        for cell in row:
            cell.draw(win)

# --- MAIN ---
def main():
    global ROWS, CELL_SIZE
    pygame.init()
    win = pygame.display.set_mode((WIDTH, WIDTH))
    pygame.display.set_caption("A*")

    clock = pygame.time.Clock()

    def create_maze(rows):
        global CELL_SIZE
        CELL_SIZE = WIDTH // rows
        grid = [[Cell(r, c) for c in range(rows)] for r in range(rows)]
        generate_maze(grid, grid[0][0])
        return grid

    grid = create_maze(ROWS)
    start = grid[0][0]
    end = grid[ROWS-1][ROWS-1]
    path = []

    running = True
    while running:
        clock.tick(FPS)
        draw_grid(win, grid)
        start.highlight(win, BLUE)
        end.highlight(win, RED)
        for cell in path:
            cell.highlight(win, YELLOW)
        pygame.display.update()

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_SPACE:
                    path = astar(grid, start, end, win)
                if event.key == pygame.K_UP:
                    ROWS += 1
                    grid = create_maze(ROWS)
                    start = grid[0][0]
                    end = grid[ROWS-1][ROWS-1]
                    path = []
                if event.key == pygame.K_DOWN and ROWS > 5:
                    ROWS -= 1
                    grid = create_maze(ROWS)
                    start = grid[0][0]
                    end = grid[ROWS-1][ROWS-1]
                    path = []

    pygame.quit()

if __name__ == "__main__":
    main()