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 (whether they are in a vehicle or at a location) and the road network to compute the minimal number of actions.

    # Assumptions
    - Packages can be at a location or inside a vehicle.
    - Vehicles can move between connected locations via roads.
    - The capacity of vehicles is respected when picking up or dropping packages.

    # Heuristic Initialization
    - Extract the goal locations for each package from the task goals.
    - Extract the road network from the static facts to compute distances between locations.
    - Extract the capacity relationships between sizes from the static facts.

    # Step-By-Step Thinking for Computing Heuristic
    1. For each package, determine its current location and whether it is inside a vehicle.
    2. If the package is already at its goal location, no action is required.
    3. If the package is inside a vehicle, determine the vehicle's current location.
    4. Compute the minimal number of actions required to move the package to its goal location:
       - If the package is not in a vehicle, it must be picked up by a vehicle.
       - The vehicle must drive to the package's current location, pick it up, drive to the goal location, and drop it.
    5. Sum the actions required for all packages to get the total heuristic value.
    """

    def __init__(self, task):
        """
        Initialize the heuristic by extracting:
        - Goal locations for each package.
        - Road network from static facts.
        - Capacity relationships from static facts.
        """
        self.goals = task.goals  # Goal conditions.
        static_facts = task.static  # Facts that are not affected by actions.

        # Map locations to their connected locations using "road" relationships.
        self.road_network = {}
        for fact in static_facts:
            if match(fact, "road", "*", "*"):
                l1, l2 = get_parts(fact)[1], get_parts(fact)[2]
                if l1 not in self.road_network:
                    self.road_network[l1] = set()
                if l2 not in self.road_network:
                    self.road_network[l2] = set()
                self.road_network[l1].add(l2)
                self.road_network[l2].add(l1)

        # Store 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

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

        # Track where packages and vehicles are currently located.
        current_locations = {}
        for fact in state:
            predicate, *args = get_parts(fact)
            if predicate in ["at", "in"]:  # Track both direct location and inside vehicle.
                obj, location = args
                current_locations[obj] = location

        total_cost = 0  # Initialize action cost counter.

        for package, goal_location in self.goal_locations.items():
            # Get the current location of the package (could be a location or a vehicle).
            current_location = current_locations[package]

            # Check if the package is inside a vehicle.
            in_vehicle = current_location.startswith("v")

            if in_vehicle:
                # Retrieve the physical location of the vehicle.
                current_location = current_locations[current_location]

            # If the package is already at its goal location, no action is required.
            if current_location == goal_location:
                continue

            # Compute the minimal number of actions required to move the package to its goal location.
            # 1. Drive to the package's current location (if not already there).
            # 2. Pick up the package.
            # 3. Drive to the goal location.
            # 4. Drop the package.
            total_cost += 4  # Each step costs 1 action.

        return total_cost
