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., "(on b1 b2)".
    - `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 BlocksworldHeuristic(Heuristic):
    """
    A domain-dependent heuristic for the Blocksworld domain.

    # Summary
    This heuristic estimates the number of actions required to reach the goal state by counting the number of blocks that are not in their correct position or not in the correct stack. It also considers the number of blocks that need to be moved to clear the way for other blocks.

    # Assumptions:
    - The goal state specifies the desired configuration of blocks, including their positions on the table and their stacking order.
    - The heuristic assumes that each block can be moved independently, and the cost of moving a block is proportional to the number of blocks on top of it.

    # Heuristic Initialization
    - Extract the goal configuration of blocks, including their positions on the table and their stacking order.
    - Identify the current state of the blocks, including their positions and stacking order.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current state of each block, including whether it is on the table, on another block, or being held.
    2. Compare the current state of each block with its goal state.
    3. For each block that is not in its goal position:
       - If the block is on the table but should be on another block, add a cost for moving it.
       - If the block is on another block but should be on the table, add a cost for moving it.
       - If the block is on the wrong block, add a cost for moving it to the correct position.
    4. For each block that is in the correct position but has blocks on top of it that are not in their correct positions, add a cost for moving those blocks.
    5. Sum the costs for all blocks to get the total heuristic value.
    """

    def __init__(self, task):
        """Initialize the heuristic by extracting goal conditions and static facts."""
        self.goals = task.goals  # Goal conditions.
        self.static = task.static  # Static facts.

        # Extract goal configuration of blocks.
        self.goal_config = {}
        for goal in self.goals:
            predicate, *args = get_parts(goal)
            if predicate == "on":
                block, under_block = args
                self.goal_config[block] = under_block
            elif predicate == "on-table":
                block = args[0]
                self.goal_config[block] = "table"

    def __call__(self, node):
        """Estimate the number of actions required to reach the goal state."""
        state = node.state  # Current world state.

        # Track the current configuration of blocks.
        current_config = {}
        for fact in state:
            predicate, *args = get_parts(fact)
            if predicate == "on":
                block, under_block = args
                current_config[block] = under_block
            elif predicate == "on-table":
                block = args[0]
                current_config[block] = "table"

        total_cost = 0  # Initialize the heuristic cost.

        # Compare current configuration with goal configuration.
        for block, goal_under_block in self.goal_config.items():
            current_under_block = current_config.get(block, None)

            if current_under_block != goal_under_block:
                # Block is not in the correct position.
                total_cost += 1

                # If the block is on another block, add a cost for moving it.
                if current_under_block != "table":
                    total_cost += 1

                # If the block is on the wrong block, add a cost for moving it to the correct position.
                if goal_under_block != "table" and current_under_block != goal_under_block:
                    total_cost += 1

        return total_cost
