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 their color, but this requires an action.
    - 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 information about tile adjacencies (up, down, left, right).
    - Extract the available colors.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the tiles that still need to be painted (i.e., tiles that are either clear or painted with the wrong color).
    2. For each robot, calculate 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 color.
    4. Sum the distances and color changes to get the total heuristic value.
    """

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

        # Map tiles to their goal colors.
        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 tile adjacencies.
        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))

        # Extract available colors.
        self.available_colors = {
            get_parts(fact)[1] for fact in static_facts if match(fact, "available-color", "*")
        }

    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():
            current_color = None
            for fact in state:
                if match(fact, "painted", tile, "*"):
                    current_color = get_parts(fact)[2]
                    break
            if current_color != goal_color:
                unpainted_tiles.append(tile)

        if not unpainted_tiles:
            return 0  # Goal state reached.

        # Calculate the distance of each robot to the nearest unpainted tile.
        robot_positions = {}
        for fact in state:
            if match(fact, "robot-at", "*", "*"):
                robot, tile = get_parts(fact)[1], get_parts(fact)[2]
                robot_positions[robot] = tile

        total_cost = 0
        for tile in unpainted_tiles:
            min_distance = float('inf')
            for robot, position in robot_positions.items():
                distance = self.calculate_distance(position, tile)
                if distance < min_distance:
                    min_distance = distance
            total_cost += min_distance

        # Estimate the number of color changes required.
        for robot in robot_positions:
            current_color = None
            for fact in state:
                if match(fact, "robot-has", robot, "*"):
                    current_color = get_parts(fact)[2]
                    break
            if current_color is None:
                continue
            for tile in unpainted_tiles:
                if self.goal_colors[tile] != current_color:
                    total_cost += 1  # One color change per robot per tile.

        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`.
        This is a simple BFS-based distance calculation.
        """
        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
            if current_tile in visited:
                continue
            visited.add(current_tile)
            for _, neighbor in self.adjacent.get(current_tile, []):
                queue.append((neighbor, distance + 1))

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