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 distance of robots to these tiles, and the
    number of color changes required.

    # Assumptions
    - Robots can move freely between adjacent tiles.
    - Robots can change colors, but this incurs an additional action cost.
    - The heuristic assumes that robots can paint tiles in the most efficient order.

    # Heuristic Initialization
    - Extract the goal conditions for each tile (i.e., the required color).
    - Extract the static facts (e.g., adjacency relationships between tiles) to compute distances.
    - Initialize data structures to store the goal colors and adjacency relationships.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the tiles that still need to be painted to match the goal conditions.
    2. For each robot, compute the distance to the nearest unpainted tile.
    3. Estimate the number of color changes required for each robot to paint the tiles with the correct colors.
    4. Sum the distances and color change 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 tiles to their required colors 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.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] = []
                self.adjacency[tile1].append(tile2)

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

        # Identify tiles that still need to be painted.
        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)

        if not unpainted_tiles:
            return 0  # Goal state reached.

        # Find the current location of each robot.
        robot_locations = {}
        for fact in state:
            if match(fact, "robot-at", "*", "*"):
                robot, tile = get_parts(fact)[1:]
                robot_locations[robot] = tile

        # Compute the distance from each robot to the nearest unpainted tile.
        total_distance = 0
        for tile in unpainted_tiles:
            min_distance = float('inf')
            for robot, location in robot_locations.items():
                distance = self._compute_distance(location, tile)
                if distance < min_distance:
                    min_distance = distance
            total_distance += min_distance

        # Estimate the number of color changes required.
        color_changes = 0
        for robot, location in robot_locations.items():
            current_color = None
            for fact in state:
                if match(fact, "robot-has", robot, "*"):
                    current_color = get_parts(fact)[2]
                    break
            if current_color:
                for tile in unpainted_tiles:
                    if self.goal_colors[tile] != current_color:
                        color_changes += 1
                        break

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

    def _compute_distance(self, start, goal):
        """
        Compute the Manhattan distance between two tiles using adjacency relationships.
        """
        # Extract coordinates from tile names (assuming format tile_X_Y).
        x1, y1 = map(int, start.split('_')[1:])
        x2, y2 = map(int, goal.split('_')[1:])
        return abs(x1 - x2) + abs(y1 - y2)
