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 need to clear blocks that are in the way of moving other blocks.

    # Assumptions:
    - The goal state specifies the desired stacking of blocks.
    - Blocks can only be moved one at a time.
    - The arm can hold only one block at a time.

    # Heuristic Initialization
    - Extract the goal conditions to determine the desired stacking of blocks.
    - Identify the current state of the blocks (on-table, on, clear, etc.).

    # 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 (for the action needed to move it).
    3. If a block is on top of another block that is not its goal position, increment the heuristic value by 1 (for the action needed to unstack it).
    4. If a block is being held by the arm, increment the heuristic value by 1 (for the action needed to put it down or stack it).
    5. If a block is clear but needs to be stacked under another block, increment the heuristic value by 1 (for the action needed to stack it).
    """

    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 goal stacking relationships.
        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] = None  # None indicates on-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 stacking relationships.
        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] = None  # None indicates on-table.

        # Initialize the heuristic value.
        heuristic_value = 0

        # Check each block's current position against its goal position.
        for block, goal_under_block in self.goal_stacking.items():
            current_under_block = current_stacking.get(block, None)

            # If the block is not in its correct position, increment the heuristic.
            if current_under_block != goal_under_block:
                heuristic_value += 1

            # If the block is on top of another block that is not its goal position, increment the heuristic.
            if current_under_block is not None and current_under_block != goal_under_block:
                heuristic_value += 1

        # Check if the arm is holding a block.
        if "(holding *)" in state:
            heuristic_value += 1

        # Check if any block is clear but needs to be stacked under another block.
        for fact in state:
            if match(fact, "clear", "*"):
                block = get_parts(fact)[1]
                if block in self.goal_stacking and self.goal_stacking[block] is not None:
                    heuristic_value += 1

        return heuristic_value
