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., "(painted tile_1_1 white)".
    - `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 current state of the tiles, the robots' positions, and the colors they are carrying.

    # Assumptions
    - Robots can move between adjacent tiles (up, down, left, right).
    - Robots can change their color if needed.
    - Each tile must be painted with the correct color as specified in the goal.

    # Heuristic Initialization
    - Extract the goal conditions for each tile.
    - Extract the static facts (e.g., adjacency relationships between tiles) to build a graph representation of the floor.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the tiles that are not yet painted with the correct color.
    2. For each unpainted tile, calculate the minimum number of moves required for a robot to reach it.
    3. If the robot needs to change its color to paint the tile, add an additional action for the color change.
    4. Sum the total number of moves and color changes required to paint all tiles correctly.
    """

    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 adjacent tiles.
        self.adjacency = {}
        for fact in static_facts:
            if match(fact, "up", "*", "*"):
                _, tile1, tile2 = get_parts(fact)
                self.adjacency.setdefault(tile1, []).append(tile2)
            elif match(fact, "down", "*", "*"):
                _, tile1, tile2 = get_parts(fact)
                self.adjacency.setdefault(tile1, []).append(tile2)
            elif match(fact, "left", "*", "*"):
                _, tile1, tile2 = get_parts(fact)
                self.adjacency.setdefault(tile1, []).append(tile2)
            elif match(fact, "right", "*", "*"):
                _, tile1, tile2 = get_parts(fact)
                self.adjacency.setdefault(tile1, []).append(tile2)

        # Store goal conditions for each tile.
        self.goal_colors = {}
        for goal in self.goals:
            if match(goal, "painted", "*", "*"):
                _, tile, color = get_parts(goal)
                self.goal_colors[tile] = color

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

        # Track the current color of each robot.
        robot_colors = {}
        for fact in state:
            if match(fact, "robot-has", "*", "*"):
                _, robot, color = get_parts(fact)
                robot_colors[robot] = color

        # Track the current position of each robot.
        robot_positions = {}
        for fact in state:
            if match(fact, "robot-at", "*", "*"):
                _, robot, tile = get_parts(fact)
                robot_positions[robot] = tile

        # Track the current color of each tile.
        tile_colors = {}
        for fact in state:
            if match(fact, "painted", "*", "*"):
                _, tile, color = get_parts(fact)
                tile_colors[tile] = color

        total_cost = 0  # Initialize action cost counter.

        for tile, goal_color in self.goal_colors.items():
            current_color = tile_colors.get(tile, None)
            if current_color != goal_color:
                # Tile needs to be painted with the correct color.
                # Find the closest robot that can paint this tile.
                min_distance = float('inf')
                best_robot = None
                for robot, position in robot_positions.items():
                    distance = self._calculate_distance(position, tile)
                    if distance < min_distance:
                        min_distance = distance
                        best_robot = robot

                # Add the number of moves required to reach the tile.
                total_cost += min_distance

                # Check if the robot needs to change its color.
                if robot_colors[best_robot] != goal_color:
                    total_cost += 1  # Add a color change action.

                # Add the painting action.
                total_cost += 1

        return total_cost

    def _calculate_distance(self, start_tile, goal_tile):
        """
        Calculate the minimum number of moves required to go from `start_tile` to `goal_tile`.
        Uses a simple breadth-first search (BFS) 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, distance = queue.popleft()
            if current_tile == goal_tile:
                return distance
            visited.add(current_tile)
            for neighbor in self.adjacency.get(current_tile, []):
                if neighbor not in visited:
                    queue.append((neighbor, distance + 1))

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