import math
from typing import Callable
import re
from graph import Graph, AdjacencyListGraph, AdjacencyMatrixGraph, NodeColor, Vertex
import heapq

def a_star(self, start_name: str, end_name: str, heuristic: Callable[[Vertex], float]):
    color_map = {}                  # maps vertices to their color
    distance_map = {}               # maps vertices to their distance from the start vertex
    predecessor_map = {}            # maps vertices to their predecessor in the traversal tree

    def cost(vertex: Vertex) -> float:
        """Compute the cost of the path to the given vertex."""
        if distance_map[vertex] is None:
            return math.inf
        return distance_map[vertex] + heuristic(vertex)

    Vertex.__lt__ = lambda self, other: cost(self) < cost(other)

    # Initialize the maps
    for vertex in self.all_vertices():
        color_map[vertex] = NodeColor.WHITE
        distance_map[vertex] = None
        predecessor_map[vertex] = None

    # Start at the given vertex
    start_node = self.get_vertex(start_name)
    color_map[start_node] = NodeColor.GRAY
    distance_map[start_node] = 0

    # Initialize the queue with the start vertex
    queue = [] # Use a list instead of a deque, because we need to sort it with
    heapq.heappush(queue,  start_node)

    # Process the queue
    while len(queue) > 0:
        vertex = heapq.heappop(queue)
        if color_map[vertex] == NodeColor.BLACK:
            # Skip already processed vertices (possibly with a higher cost)
            continue
        if vertex.value == end_name:
            # Return the distance and predecessor maps
            return distance_map, predecessor_map
        for dest, weight in self.get_adjacent_vertices_with_weight(vertex.value):
            if color_map[dest] == NodeColor.BLACK:
                continue
            f = distance_map[vertex] + weight + heuristic(dest)
            if color_map[dest] == NodeColor.GRAY and f > cost(dest):
                continue
            predecessor_map[dest] = vertex
            distance_map[dest] = distance_map[vertex] + weight
            heapq.heappush(queue, dest)
            if color_map[dest] == NodeColor.WHITE:
                color_map[dest] = NodeColor.GRAY
        color_map[vertex] = NodeColor.BLACK

    # Return the distance and predecessor maps if no path was found
    return None, None


# Add the a_star method to the Graph classes
AdjacencyListGraph.a_star = a_star
AdjacencyMatrixGraph.a_star = a_star


if __name__ == "__main__":

    def read_labyrinth_into_graph(graph: Graph, filename: str):
        """Read a labyrinth from a file into a graph. The file format is a grid of characters:"""
        start = None
        end = None
        with open(filename, "r") as file:
            nodes = []
            lines = file.readlines()
            for y, line in enumerate(lines):
                for x, char in enumerate(line):
                    if char in ' AS':
                        name = pos_to_nodename((x, y))
                        graph.insert_vertex(name)
                        nodes.append((x, y))
                    if char == 'A':
                        end = (x, y)
                    if char == 'S':
                        start = (x, y)
            for x, y in nodes:
                name1 = f"x{x}y{y}"
                for neighbor in [(x - 1, y), (x, y - 1), (x, y + 1), (x + 1, y)]:
                    if neighbor in nodes:
                        name_neighbor = pos_to_nodename(neighbor)
                        graph.connect(name1, name_neighbor, 1)
        return start, end, lines


    def nodename_to_pos(nodename):
        """Convert a node name to a position (x, y)."""
        m = re.match(r"(^x(\d*)y(\d*))", nodename)
        if m:
            return (int(m.group(2)), int(m.group(3)))
        return None


    def pos_to_nodename(pos):
        """Convert a position (x, y) to a node name."""
        x, y = pos
        return f"x{x}y{y}"


    def get_heuristic(end) -> Callable[[Vertex], float]:
        """Return a heuristic function for the given end position."""
        def heuristic(v: Vertex) -> float:
            x, y = nodename_to_pos(v.value)
            x1, y1 = end
            # Euclidean distance
            return math.sqrt(abs(x - x1)**2 + abs(y - y1)**2)

        return heuristic

    def update_lines(lines, distance_map, path):
        """Update the lines with the path found by the A* algorithm."""
        def replace_at(x, y, replacement):
            if lines[y][x] not in " .":
                return
            lines[y] = lines[y][:x] + replacement + lines[y][x+1:]

        for node in distance_map.keys():
            if distance_map[node] is not None:
                x, y = nodename_to_pos(node.value)
                replace_at(x, y, ".")
        for node in path:
            x, y = nodename_to_pos(node)
            replace_at(x, y, "*")
        return lines

    graph = AdjacencyListGraph()
    #graph = AdjacencyMatrixGraph()
    start, end, lines = read_labyrinth_into_graph(graph, "../../labyrinth.txt")
    distance_map, predecessor_map = graph.a_star(pos_to_nodename(start), pos_to_nodename(end), get_heuristic(end))
    endname = pos_to_nodename(end)
    lines = update_lines(lines, distance_map, graph.path(endname, predecessor_map))
    for line in lines:
        print(line, end="")