from fnmatch import fnmatch
from heuristics.heuristic_base import Heuristic


def get_parts(fact):
    """Extract the components of a PDDL fact by removing parentheses and splitting the string."""
    return fact[1:-1].split()


def match(fact, *args):
    """
    Check if a PDDL fact matches a given pattern.

    - `fact`: The complete fact as a string, e.g., "(robot-at robot1 tile_0_1)".
    - `args`: The expected pattern (wildcards `*` allowed).
    - Returns `True` if the fact matches the pattern, else `False`.
    """
    parts = get_parts(fact)
    return all(fnmatch(part, arg) for part, arg in zip(parts, args))


class FloortileHeuristic(Heuristic):
    """
    A domain-dependent heuristic for the Floortile domain.

    # Summary
    This heuristic estimates the number of actions required to paint all tiles according to the goal configuration.
    It considers the following:
    - The number of tiles that still need to be painted.
    - The distance between the robot's current position and the tiles that need to be painted.
    - The number of color changes required for the robot.

    # Assumptions
    - The robot can only paint adjacent tiles (up, down, left, right).
    - The robot must move to a tile before painting it.
    - The robot may need to change its color before painting a tile.

    # Heuristic Initialization
    - Extract the goal conditions for each tile (i.e., which color each tile should be painted).
    - Extract the static facts (e.g., adjacency relationships between tiles) to compute distances.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the tiles that are not yet painted in the goal color.
    2. For each unpainted tile, compute the Manhattan distance from the robot's current position.
    3. Sum the distances to estimate the number of moves required.
    4. Add the number of color changes required (one per unique color change needed).
    5. The total heuristic value is the sum of the move and color change costs.
    """

    def __init__(self, task):
        """
        Initialize the heuristic by extracting:
        - Goal conditions for each tile.
        - Static facts (adjacency relationships between tiles).
        """
        self.goals = task.goals  # Goal conditions.
        static_facts = task.static  # Facts that are not affected by actions.

        # Map each tile to its goal color.
        self.goal_colors = {}
        for goal in self.goals:
            predicate, *args = get_parts(goal)
            if predicate == "painted":
                tile, color = args
                self.goal_colors[tile] = color

        # Build adjacency graph for tiles.
        self.adjacency = {}
        for fact in static_facts:
            predicate, *args = get_parts(fact)
            if predicate in ["up", "down", "left", "right"]:
                tile1, tile2 = args
                if tile1 not in self.adjacency:
                    self.adjacency[tile1] = set()
                self.adjacency[tile1].add(tile2)

    def __call__(self, node):
        """Compute an estimate of the minimal number of required actions."""
        state = node.state  # Current world state.

        # Identify the robot's current position and color.
        robot_position = None
        robot_color = None
        for fact in state:
            predicate, *args = get_parts(fact)
            if predicate == "robot-at":
                robot, tile = args
                robot_position = tile
            elif predicate == "robot-has":
                robot, color = args
                robot_color = color

        # Count the number of tiles that are not yet painted in the goal color.
        unpainted_tiles = []
        for tile, goal_color in self.goal_colors.items():
            if not any(match(fact, "painted", tile, goal_color) for fact in state):
                unpainted_tiles.append(tile)

        # Compute the Manhattan distance from the robot to each unpainted tile.
        total_distance = 0
        for tile in unpainted_tiles:
            # Compute the distance using BFS (since adjacency is not grid-based).
            distance = self._bfs_distance(robot_position, tile)
            total_distance += distance

        # Count the number of color changes required.
        color_changes = 0
        for tile in unpainted_tiles:
            goal_color = self.goal_colors[tile]
            if goal_color != robot_color:
                color_changes += 1

        # The heuristic value is the sum of distances and color changes.
        return total_distance + color_changes

    def _bfs_distance(self, start, goal):
        """
        Compute the shortest path distance between two tiles using BFS.

        - `start`: The starting tile.
        - `goal`: The goal tile.
        - Returns the number of moves required to reach the goal tile.
        """
        if start == goal:
            return 0

        visited = set()
        queue = [(start, 0)]

        while queue:
            current, distance = queue.pop(0)
            if current == goal:
                return distance
            visited.add(current)
            for neighbor in self.adjacency.get(current, []):
                if neighbor not in visited:
                    queue.append((neighbor, distance + 1))

        return float("inf")  # If no path exists (should not happen in valid instances).
