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 conditions.
    It considers the following factors:
    - 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 to paint the tiles with the correct colors.

    # Assumptions
    - The robot can only paint adjacent tiles (up, down, left, right).
    - The robot must change its color if it needs to paint a tile with a different color than the one it currently holds.
    - The robot can move freely between tiles as long as they are clear.

    # 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 determine the movement and painting possibilities.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the tiles that still need to be painted according to the goal conditions.
    2. For each unpainted tile, calculate the minimum number of moves required for the robot to reach it from its current position.
    3. If the robot needs to change its color to paint a tile, add an additional action for the color change.
    4. Sum the total number of moves and color changes to estimate the total number of actions required.
    """

    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

        # Extract adjacency relationships between tiles.
        self.adjacent = {}
        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.adjacent:
                    self.adjacent[tile1] = []
                self.adjacent[tile1].append((predicate, tile2))

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

        # Track 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

        total_cost = 0  # Initialize action cost counter.

        # Iterate over all tiles that need to be painted.
        for tile, goal_color in self.goal_colors.items():
            # Check if the tile is already painted with the correct color.
            if f"(painted {tile} {goal_color})" in state:
                continue

            # Calculate the minimum number of moves required to reach the tile.
            if robot_position == tile:
                move_cost = 0
            else:
                move_cost = self._calculate_move_cost(robot_position, tile)

            # Check if a color change is needed.
            if robot_color != goal_color:
                color_change_cost = 1
            else:
                color_change_cost = 0

            # Add the cost of moving and painting the tile.
            total_cost += move_cost + color_change_cost + 1  # +1 for the paint action.

        return total_cost

    def _calculate_move_cost(self, start_tile, goal_tile):
        """
        Calculate the minimum number of moves required to go from `start_tile` to `goal_tile`.
        This is a simple BFS-based implementation to find the shortest path.
        """
        from collections import deque

        if start_tile == goal_tile:
            return 0

        visited = set()
        queue = deque([(start_tile, 0)])

        while queue:
            current_tile, cost = queue.popleft()
            if current_tile == goal_tile:
                return cost

            if current_tile in visited:
                continue
            visited.add(current_tile)

            for direction, neighbor in self.adjacent.get(current_tile, []):
                queue.append((neighbor, cost + 1))

        return float('inf')  # If no path is found (should not happen in valid instances).
