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 correctly stacked. It also considers the number of blocks that need to be moved to achieve the correct stacking.

    # Assumptions:
    - The goal state specifies the desired stacking of blocks.
    - Blocks can only be moved one at a time.
    - The heuristic does not need to be admissible, so it can overestimate the number of actions.

    # Heuristic Initialization
    - Extract the goal conditions to determine the desired stacking of blocks.
    - Identify the correct positions and stacking relationships for each block.

    # Step-By-Step Thinking for Computing Heuristic
    1. For each block, check if it is in its correct position as specified by the goal.
    2. If a block is not in its correct position, increment the heuristic value by 1.
    3. For each block that is not correctly stacked (i.e., not on the correct block or not on the table as specified by the goal), increment the heuristic value by 1.
    4. If a block is being held by the arm, it will require at least one action to place it correctly, so increment the heuristic value by 1.
    5. The total heuristic value is the sum of all these increments.
    """

    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 (not used in this heuristic).

        # Extract the desired stacking from the goal conditions.
        self.goal_stacking = {}
        for goal in self.goals:
            predicate, *args = get_parts(goal)
            if predicate == "on":
                block, under_block = args
                self.goal_stacking[block] = under_block
            elif predicate == "on-table":
                block = args[0]
                self.goal_stacking[block] = "table"

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

        heuristic_value = 0  # Initialize the heuristic value.

        # Check if the goal is already reached.
        if self.goals <= state:
            return 0

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

        # Compare current stacking with goal stacking.
        for block, goal_under_block in self.goal_stacking.items():
            if block not in current_stacking:
                # Block is being held by the arm.
                heuristic_value += 1
            else:
                current_under_block = current_stacking[block]
                if current_under_block != goal_under_block:
                    # Block is not in the correct position.
                    heuristic_value += 1

        return heuristic_value
