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., "(at bob shed)".
    - `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 SpannerHeuristic(Heuristic):
    """
    A domain-dependent heuristic for the Spanner domain.

    # Summary
    This heuristic estimates the number of actions needed to tighten all loose nuts in the Spanner domain.
    It considers the following:
    - The man must pick up a usable spanner if not already carrying one.
    - The man must walk to the location of each loose nut.
    - The man must tighten each loose nut using a usable spanner.

    # Assumptions:
    - The man can carry only one spanner at a time.
    - A spanner becomes unusable after tightening a nut.
    - The man must walk to the location of each nut to tighten it.
    - The heuristic does not account for the optimal path between locations.

    # Heuristic Initialization
    - Extract the goal conditions (tightened nuts) and static facts (links between locations).
    - Create a mapping of locations to their connected locations using the static facts.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the number of loose nuts that need to be tightened.
    2. Check if the man is carrying a usable spanner:
       - If not, add the cost of picking up a usable spanner.
    3. For each loose nut:
       - Add the cost of walking to the nut's location.
       - Add the cost of tightening the nut (using the spanner).
    4. If the man is not at the starting location of the path, add the cost of walking to the starting location.
    """

    def __init__(self, task):
        """Initialize the heuristic by extracting goal conditions and static facts."""
        self.goals = task.goals  # Goal conditions.
        static_facts = task.static  # Facts that are not affected by actions.

        # Map locations to their connected locations using "link" relationships.
        self.location_links = {}
        for fact in static_facts:
            if match(fact, "link", "*", "*"):
                _, loc1, loc2 = get_parts(fact)
                if loc1 not in self.location_links:
                    self.location_links[loc1] = set()
                if loc2 not in self.location_links:
                    self.location_links[loc2] = set()
                self.location_links[loc1].add(loc2)
                self.location_links[loc2].add(loc1)

    def __call__(self, node):
        """Estimate the minimum cost to tighten all loose nuts."""
        state = node.state  # Current world state.

        # Count the number of loose nuts that need to be tightened.
        loose_nuts = sum(1 for fact in state if match(fact, "loose", "*"))

        # Check if the man is carrying a usable spanner.
        carrying_spanner = any(
            match(fact, "carrying", "*", "*") and match(fact, "*", "*", "spanner*")
            for fact in state
        )
        usable_spanner = any(match(fact, "usable", "*") for fact in state)

        # Initialize the heuristic cost.
        total_cost = 0

        # If the man is not carrying a usable spanner, add the cost of picking one up.
        if not (carrying_spanner and usable_spanner):
            total_cost += 1  # Cost of picking up a spanner.

        # Add the cost of walking to each loose nut and tightening it.
        total_cost += loose_nuts * 2  # 1 for walking, 1 for tightening.

        return total_cost
