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 pick up a spanner and walk to the nut's location in the most efficient way.

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

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current location of the man and the spanners.
    2. Identify the locations of all loose nuts.
    3. If the man is not carrying a usable spanner:
       - Find the nearest usable spanner and compute the walking distance to pick it up.
       - Add the cost of picking up the spanner.
    4. For each loose nut:
       - Compute the walking distance from the man's current location (or the spanner's location) to the nut's location.
       - Add the cost of walking and tightening the nut.
    5. Sum the total cost of all actions required to tighten all loose nuts.
    """

    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

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

        # Extract the locations of all loose nuts.
        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, "*"):
                        loose_nuts.append((nut, get_parts(loc_fact)[2]))
                        break

        # Extract the locations of all usable spanners.
        usable_spanners = []
        for fact in state:
            if match(fact, "usable", "*"):
                spanner = get_parts(fact)[1]
                for loc_fact in state:
                    if match(loc_fact, "at", spanner, "*"):
                        usable_spanners.append((spanner, get_parts(loc_fact)[2]))
                        break

        # Check if the man is carrying a usable spanner.
        carrying_spanner = None
        for fact in state:
            if match(fact, "carrying", "bob", "*"):
                carrying_spanner = get_parts(fact)[2]
                break

        total_cost = 0  # Initialize the heuristic cost.

        # If the man is not carrying a usable spanner, find the nearest one.
        if carrying_spanner is None:
            if not usable_spanners:
                return float('inf')  # No usable spanners available.
            # Find the nearest usable spanner.
            nearest_spanner_location = None
            min_distance = float('inf')
            for spanner, location in usable_spanners:
                distance = self._compute_distance(man_location, location)
                if distance < min_distance:
                    min_distance = distance
                    nearest_spanner_location = location
            total_cost += min_distance  # Walk to the spanner.
            total_cost += 1  # Pick up the spanner.
            man_location = nearest_spanner_location  # Update man's location.

        # For each loose nut, compute the cost to tighten it.
        for nut, nut_location in loose_nuts:
            distance = self._compute_distance(man_location, nut_location)
            total_cost += distance  # Walk to the nut.
            total_cost += 1  # Tighten the nut.
            man_location = nut_location  # Update man's location.

        return total_cost

    def _compute_distance(self, start, end):
        """
        Compute the shortest path distance between two locations using BFS.

        - `start`: The starting location.
        - `end`: The target location.
        - Returns the number of steps (actions) required to move from `start` to `end`.
        """
        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.
