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 number of tiles that still need to be painted, the current color of the robot, and the
    distance the robot needs to travel to reach unpainted tiles.

    # Assumptions
    - The robot can only paint tiles adjacent to its current position.
    - The robot must change color if it needs to paint a tile with a different color than the one it currently holds.
    - The robot can move freely between adjacent tiles, provided they are clear.

    # Heuristic Initialization
    - Extract the goal conditions for each tile (i.e., the required color for each tile).
    - 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 still need to be painted by comparing the current state with the goal conditions.
    2. For each unpainted tile, compute the Manhattan distance from the robot's current position to the tile.
    3. If the robot needs to change color to paint a tile, add an additional action cost for the color change.
    4. Sum the distances and additional costs 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 required color in the goal state.
        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(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

        # Identify tiles that still need to be painted.
        unpainted_tiles = []
        for tile, goal_color in self.goal_colors.items():
            if f"(painted {tile} {goal_color})" not in state:
                unpainted_tiles.append(tile)

        # If all tiles are painted, the heuristic is 0.
        if not unpainted_tiles:
            return 0

        # Compute the Manhattan distance between two tiles.
        def manhattan_distance(tile1, tile2):
            x1, y1 = map(int, tile1.split('_')[1:])
            x2, y2 = map(int, tile2.split('_')[1:])
            return abs(x1 - x2) + abs(y1 - y2)

        # Estimate the total cost.
        total_cost = 0
        for tile in unpainted_tiles:
            # Distance to reach the tile.
            distance = manhattan_distance(robot_position, tile)
            total_cost += distance

            # Check if the robot needs to change color.
            goal_color = self.goal_colors[tile]
            if robot_color != goal_color:
                total_cost += 1  # Cost for changing color.

        return total_cost
