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 arm can hold only one block at a time.
    - Blocks must be clear to be picked up or stacked.

    # Heuristic Initialization
    - Extract the goal conditions for each block, including their positions and stacking relationships.
    - Identify the initial state of the blocks, including their current positions and stacking relationships.

    # Step-By-Step Thinking for Computing Heuristic
    1. For each block, check if it is in its correct position as specified in the goal state.
    2. If a block is not in its correct position, increment the heuristic value by 1 (since it needs to be moved).
    3. If a block is in its correct position but the block below it is not, increment the heuristic value by 1 (since the block below needs to be moved first).
    4. If a block is being held by the arm, increment the heuristic value by 1 (since it needs to be placed).
    5. If the arm is empty but a block needs to be picked up, increment the heuristic value by 1 (since the arm needs to pick up the block).
    """

    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 needed for this heuristic).

        # Extract goal relationships for each block.
        self.goal_on = {}
        self.goal_on_table = set()
        for goal in self.goals:
            predicate, *args = get_parts(goal)
            if predicate == "on":
                block, under_block = args
                self.goal_on[block] = under_block
            elif predicate == "on-table":
                block = args[0]
                self.goal_on_table.add(block)

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

        # Track the current positions and stacking of blocks.
        current_on = {}
        current_on_table = set()
        holding = None
        for fact in state:
            predicate, *args = get_parts(fact)
            if predicate == "on":
                block, under_block = args
                current_on[block] = under_block
            elif predicate == "on-table":
                block = args[0]
                current_on_table.add(block)
            elif predicate == "holding":
                holding = args[0]

        total_cost = 0  # Initialize the heuristic cost.

        # Check each block's position against the goal.
        for block in self.goal_on_table | set(self.goal_on.keys()):
            if block in self.goal_on_table:
                # Block should be on the table.
                if block not in current_on_table:
                    total_cost += 1  # Needs to be moved to the table.
            else:
                # Block should be on another block.
                goal_under_block = self.goal_on[block]
                if block not in current_on or current_on[block] != goal_under_block:
                    total_cost += 1  # Needs to be moved to the correct position.

        # If the arm is holding a block, it needs to be placed.
        if holding is not None:
            total_cost += 1

        # If the arm is empty but a block needs to be picked up, increment the cost.
        if holding is None and any(block not in current_on_table for block in self.goal_on_table):
            total_cost += 1

        return total_cost
