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:
    - 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.

    # Assumptions
    - The robot can only paint adjacent tiles (up, down, left, right).
    - The robot must change color if the current color does not match the required color for the next tile.
    - 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 compute distances.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the tiles that still need to be painted (i.e., tiles that are not yet painted or are painted with the wrong color).
    2. For each such tile, compute the Manhattan distance from the robot's current position to the tile.
    3. If the robot's current color does not match the required color for the tile, add a cost for changing the color.
    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.

        # Extract adjacency relationships between tiles.
        self.adjacent = {}
        for fact in static_facts:
            if match(fact, "up", "*", "*"):
                _, tile1, tile2 = get_parts(fact)
                self.adjacent[(tile1, tile2)] = "up"
            elif match(fact, "down", "*", "*"):
                _, tile1, tile2 = get_parts(fact)
                self.adjacent[(tile1, tile2)] = "down"
            elif match(fact, "left", "*", "*"):
                _, tile1, tile2 = get_parts(fact)
                self.adjacent[(tile1, tile2)] = "left"
            elif match(fact, "right", "*", "*"):
                _, tile1, tile2 = get_parts(fact)
                self.adjacent[(tile1, tile2)] = "right"

        # Store goal conditions for each tile.
        self.goal_colors = {}
        for goal in self.goals:
            predicate, *args = get_parts(goal)
            if predicate == "painted":
                tile, color = args
                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.

        # Identify the robot's current position and color.
        robot_position = None
        robot_color = None
        for fact in state:
            if match(fact, "robot-at", "*", "*"):
                _, robot, tile = get_parts(fact)
                robot_position = tile
            elif match(fact, "robot-has", "*", "*"):
                _, robot, color = get_parts(fact)
                robot_color = color

        # 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, "*"):
                    _, _, color = get_parts(fact)
                    current_color = color
                    break
            if current_color != goal_color:
                unpainted_tiles.append((tile, goal_color))

        # Compute the heuristic value.
        total_cost = 0
        for tile, goal_color in unpainted_tiles:
            # Compute the Manhattan distance between the robot's current position and the tile.
            distance = self.compute_manhattan_distance(robot_position, tile)
            total_cost += distance

            # If the robot's current color does not match the goal color, add a cost for changing the color.
            if robot_color != goal_color:
                total_cost += 1  # Cost for changing color.

        return total_cost

    def compute_manhattan_distance(self, tile1, tile2):
        """
        Compute the Manhattan distance between two tiles based on their coordinates.
        Assumes tile names are in the format "tile_X_Y".
        """
        x1, y1 = map(int, tile1.split("_")[1:])
        x2, y2 = map(int, tile2.split("_")[1:])
        return abs(x1 - x2) + abs(y1 - y2)
