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).
    - Build a graph of locations using the static `link` facts to compute distances.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current location of the man and the spanners.
    2. Determine if the man is carrying a usable spanner.
    3. For each loose nut:
       - Compute the distance from the man's current location to the nut's location.
       - Add the cost of walking to the nut's location.
       - Add the cost of tightening the nut if a usable spanner is available.
    4. If the man is not carrying a usable spanner:
       - Add the cost of walking to the nearest usable spanner.
       - Add the cost of picking up the spanner.
    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 (tightened nuts).
        self.static = task.static  # Static facts (links between locations).

        # Build a graph of locations using the static `link` facts.
        self.location_graph = {}
        for fact in self.static:
            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
        if not man_location:
            return float('inf')  # Invalid state.

        # 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

        # Check if the carried spanner is usable.
        usable_spanner = False
        if carrying_spanner:
            for fact in state:
                if match(fact, "usable", carrying_spanner):
                    usable_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, check if goals are already satisfied.
        if not loose_nuts:
            return 0 if self.goals <= state else float('inf')

        # Compute the heuristic cost.
        total_cost = 0

        # If the man is not carrying a usable spanner, find the nearest one.
        if not usable_spanner:
            # Find all usable spanners and their locations.
            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, "*"):
                            spanner_location = get_parts(loc_fact)[2]
                            usable_spanners.append((spanner, spanner_location))
                            break

            if not usable_spanners:
                return float('inf')  # No usable spanners available.

            # Find the nearest usable spanner.
            min_distance = float('inf')
            nearest_spanner_location = None
            for spanner, spanner_location in usable_spanners:
                distance = self._compute_distance(man_location, spanner_location)
                if distance < min_distance:
                    min_distance = distance
                    nearest_spanner_location = spanner_location

            # Add the cost of walking to the nearest spanner and picking it up.
            total_cost += min_distance + 1  # 1 action to pick up the spanner.
            man_location = nearest_spanner_location  # Update man's location.

        # For each loose nut, add the cost of walking to its location and tightening it.
        for nut, nut_location in loose_nuts:
            distance = self._compute_distance(man_location, nut_location)
            total_cost += distance + 1  # 1 action to tighten the nut.
            man_location = nut_location  # Update man's location.

        return total_cost

    def _compute_distance(self, start, end):
        """
        Compute the minimum number of steps (walk actions) between two locations using BFS.

        - `start`: The starting location.
        - `end`: The target location.
        - Returns the number of steps required to walk from `start` to `end`.
        """
        if start == end:
            return 0

        visited = set()
        queue = [(start, 0)]

        while queue:
            current, steps = queue.pop(0)
            if current == end:
                return steps
            if current in visited:
                continue
            visited.add(current)
            for neighbor in self.location_graph.get(current, []):
                queue.append((neighbor, steps + 1))

        return float('inf')  # No path found.
