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 steps:
    - Walking to the location of a spanner.
    - Picking up the spanner.
    - Walking to the location of a loose nut.
    - Tightening the nut.

    # Assumptions:
    - The man can carry only one spanner at a time.
    - Each spanner can be used only once.
    - The man must walk to the location of a spanner or nut to interact with it.
    - The goal is to tighten all loose nuts.

    # Heuristic Initialization
    - Extract the goal conditions (tightened nuts) and static facts (links between locations).
    - Build a graph of locations using the static `link` facts to compute distances.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify all loose nuts that need to be tightened.
    2. For each loose nut:
        a. Find the nearest spanner that can be used to tighten it.
        b. Compute the distance from the man's current location to the spanner's location.
        c. Compute the distance from the spanner's location to the nut's location.
        d. Add the cost of picking up the spanner and tightening the nut.
    3. 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.
        static_facts = task.static  # Facts that are not affected by actions.

        # Build a graph of locations using the static `link` facts.
        self.location_graph = {}
        for fact in static_facts:
            if match(fact, "link", "*", "*"):
                _, loc1, loc2 = get_parts(fact)
                if loc1 not in self.location_graph:
                    self.location_graph[loc1] = set()
                if loc2 not in self.location_graph:
                    self.location_graph[loc2] = set()
                self.location_graph[loc1].add(loc2)
                self.location_graph[loc2].add(loc1)

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

        # Identify all loose nuts that need to be tightened.
        loose_nuts = {get_parts(fact)[1] for fact in state if match(fact, "loose", "*")}

        # If there are no loose nuts, the heuristic is 0 (goal state).
        if not loose_nuts:
            return 0

        # Identify the man's current location.
        man_location = None
        for fact in state:
            if match(fact, "at", "bob", "*"):
                man_location = get_parts(fact)[2]
                break
        if not man_location:
            return float('inf')  # Invalid state.

        # Identify all usable spanners and their locations.
        usable_spanners = {}
        for fact in state:
            if match(fact, "usable", "*"):
                spanner = get_parts(fact)[1]
                for fact2 in state:
                    if match(fact2, "at", spanner, "*"):
                        usable_spanners[spanner] = get_parts(fact2)[2]
                        break

        # Compute the heuristic value.
        total_cost = 0

        for nut in loose_nuts:
            # Find the nut's location.
            nut_location = None
            for fact in state:
                if match(fact, "at", nut, "*"):
                    nut_location = get_parts(fact)[2]
                    break
            if not nut_location:
                return float('inf')  # Invalid state.

            # Find the nearest usable spanner.
            min_cost = float('inf')
            best_spanner = None
            for spanner, spanner_location in usable_spanners.items():
                # Compute the distance from the man's location to the spanner's location.
                distance_to_spanner = self._shortest_path(man_location, spanner_location)
                # Compute the distance from the spanner's location to the nut's location.
                distance_to_nut = self._shortest_path(spanner_location, nut_location)
                # Total cost: walk to spanner, pick it up, walk to nut, tighten it.
                cost = distance_to_spanner + 1 + distance_to_nut + 1
                if cost < min_cost:
                    min_cost = cost
                    best_spanner = spanner

            if best_spanner:
                total_cost += min_cost
                # Remove the used spanner from the list of usable spanners.
                del usable_spanners[best_spanner]
            else:
                return float('inf')  # No usable spanner found.

        return total_cost

    def _shortest_path(self, start, end):
        """Compute the shortest path between two locations using BFS."""
        if start == end:
            return 0
        visited = set()
        queue = [(start, 0)]
        while queue:
            current, distance = queue.pop(0)
            if current == end:
                return distance
            if current in visited:
                continue
            visited.add(current)
            for neighbor in self.location_graph.get(current, []):
                queue.append((neighbor, distance + 1))
        return float('inf')  # No path found.
