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 p1 l1)".
    - `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 TransportHeuristic(Heuristic):
    """
    A domain-dependent heuristic for the Transport domain.

    # Summary
    This heuristic estimates the number of actions required to transport all packages to their goal locations. It considers the current state of packages, vehicles, and roads, and calculates the minimal number of actions needed to achieve the goal.

    # Assumptions
    - Packages can be at a location or inside a vehicle.
    - Vehicles can move between connected locations via roads.
    - Vehicles have limited capacity, which affects the number of packages they can carry.
    - The heuristic assumes that the most efficient path is taken, but does not guarantee admissibility.

    # Heuristic Initialization
    - Extract goal locations for each package from the task goals.
    - Extract static information about roads and capacity-predecessor relationships.
    - Build a graph of locations connected by roads for pathfinding.

    # Step-By-Step Thinking for Computing Heuristic
    1. For each package, determine its current location and goal location.
    2. If the package is already at its goal location, no action is needed.
    3. If the package is inside a vehicle, determine the vehicle's current location.
    4. Calculate the shortest path (number of road segments) from the current location to the goal location.
    5. If the package is not in a vehicle, estimate the number of actions required to pick it up and transport it to the goal.
    6. Sum the actions required for all packages to get the total heuristic value.
    """

    def __init__(self, task):
        """Initialize the heuristic by extracting goal conditions and static facts."""
        self.goals = task.goals
        self.static = task.static

        # Extract goal locations for each package
        self.goal_locations = {}
        for goal in self.goals:
            predicate, *args = get_parts(goal)
            if predicate == "at":
                package, location = args
                self.goal_locations[package] = location

        # Build a graph of locations connected by roads
        self.road_graph = {}
        for fact in self.static:
            if match(fact, "road", "*", "*"):
                _, l1, l2 = get_parts(fact)
                if l1 not in self.road_graph:
                    self.road_graph[l1] = set()
                if l2 not in self.road_graph:
                    self.road_graph[l2] = set()
                self.road_graph[l1].add(l2)
                self.road_graph[l2].add(l1)

    def __call__(self, node):
        """Compute an estimate of the minimal number of required actions."""
        state = node.state

        # Track current locations of packages and vehicles
        current_locations = {}
        for fact in state:
            predicate, *args = get_parts(fact)
            if predicate in ["at", "in"]:
                obj, location = args
                current_locations[obj] = location

        total_cost = 0

        for package, goal_location in self.goal_locations.items():
            current_location = current_locations.get(package, None)

            if current_location == goal_location:
                continue  # Package is already at its goal

            # If the package is inside a vehicle, get the vehicle's location
            if current_location.startswith("v"):
                vehicle = current_location
                current_location = current_locations.get(vehicle, None)

            # Calculate the shortest path from current_location to goal_location
            if current_location and goal_location in self.road_graph:
                path_length = self._shortest_path_length(current_location, goal_location)
                total_cost += path_length

            # Add cost for picking up and dropping the package
            total_cost += 2  # One pick-up and one drop action

        return total_cost

    def _shortest_path_length(self, start, goal):
        """Compute the shortest path length between two locations using BFS."""
        if start == goal:
            return 0

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

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

        return float('inf')  # If no path exists, return infinity
