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.
    - 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 initial state of the blocks to compare against the goal state.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current state of each block (on-table, on another block, or being held).
    2. Compare the current state of each block with its goal state.
    3. Count the number of blocks that are not in their correct position or not correctly stacked.
    4. For each block that is not in its correct position, estimate the number of actions required to move it to the correct position:
       - If a block is on the wrong block or on the table when it should be on another block, it needs to be moved.
       - If a block is being held, it needs to be placed in the correct position.
    5. Sum the estimated actions to get the heuristic value.
    """

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

        # Extract the goal state for each block.
        self.goal_state = {}
        for goal in self.goals:
            predicate, *args = get_parts(goal)
            if predicate == "on":
                block, under_block = args
                self.goal_state[block] = under_block
            elif predicate == "on-table":
                block = args[0]
                self.goal_state[block] = "table"

    def __call__(self, node):
        """Compute an estimate of the minimal number of required actions."""
        state = node.state  # Current world state.

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

        total_cost = 0  # Initialize action cost counter.

        for block, goal_under_block in self.goal_state.items():
            current_under_block = current_state.get(block, None)

            if current_under_block != goal_under_block:
                # The block is not in its correct position.
                total_cost += 1  # At least one action to move it.

                if current_under_block == "held":
                    # If the block is being held, it needs to be placed.
                    total_cost += 1
                elif current_under_block != "table":
                    # If the block is on another block, it needs to be unstacked first.
                    total_cost += 1

        return total_cost
