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 assumes that the man can always find a usable spanner if needed.

    # Heuristic Initialization
    - Extract the goal conditions (tightened nuts) and static facts (links between locations).
    - Identify all loose nuts and their locations.
    - Identify all usable spanners and their locations.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current location of the man.
    2. Check if the man is carrying a usable spanner:
       - If not, estimate the cost to pick up a usable spanner.
    3. For each loose nut:
       - Estimate the cost to walk to the nut's location.
       - Estimate the cost to tighten the nut using a usable spanner.
    4. Sum the costs for all loose nuts to get the total heuristic value.
    """

    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 (e.g., links between locations).

        # Extract all links between locations.
        self.links = {}
        for fact in self.static:
            if match(fact, "link", "*", "*"):
                parts = get_parts(fact)
                loc1, loc2 = parts[1], parts[2]
                if loc1 not in self.links:
                    self.links[loc1] = set()
                if loc2 not in self.links:
                    self.links[loc2] = set()
                self.links[loc1].add(loc2)
                self.links[loc2].add(loc1)

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

        # Identify the current location of the man.
        man_location = None
        for fact in state:
            if match(fact, "at", "bob", "*"):
                man_location = get_parts(fact)[2]
                break

        # Check if the man is carrying a usable spanner.
        carrying_spanner = False
        for fact in state:
            if match(fact, "carrying", "bob", "*"):
                carrying_spanner = True
                break

        # Identify all loose nuts and their locations.
        loose_nuts = []
        for fact in state:
            if match(fact, "loose", "*"):
                nut = get_parts(fact)[1]
                for loc_fact in state:
                    if match(loc_fact, "at", nut, "*"):
                        nut_location = get_parts(loc_fact)[2]
                        loose_nuts.append((nut, nut_location))
                        break

        # If no loose nuts, return 0 (goal state).
        if not loose_nuts:
            return 0

        # Estimate the cost to pick up a usable spanner if not already carrying one.
        cost = 0
        if not carrying_spanner:
            # Assume the man needs to walk to the nearest spanner and pick it up.
            # For simplicity, we assume a cost of 2 actions (walk + pickup).
            cost += 2

        # Estimate the cost to walk to each loose nut and tighten it.
        for nut, nut_location in loose_nuts:
            # Estimate the walking distance from the man's current location to the nut's location.
            # For simplicity, we assume a cost of 1 action per step.
            # This can be improved by calculating the actual shortest path using the links.
            if man_location != nut_location:
                cost += 1  # Walk to the nut's location.

            # Tighten the nut (1 action).
            cost += 1

        return cost
