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 considered when picking up or dropping packages.
    - The heuristic assumes that the road network is connected, and all goal locations are reachable.

    # Heuristic Initialization
    - Extract the goal locations for each package from the task goals.
    - Extract the road network from the static facts to determine connectivity between locations.
    - Extract the capacity constraints of vehicles from the static facts.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current location of each package and its goal location.
    2. Determine if the package is inside a vehicle or at a location.
    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:
       - If the package is already at its goal location, no action is needed.
       - If the package is in a vehicle, it may need to be dropped at the goal location.
       - If the package is not in a vehicle, it may need to be picked up by a vehicle and then transported.
    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 constraints of vehicles.
        """
        self.goals = task.goals  # Goal conditions.
        static_facts = task.static  # Facts that are not affected by actions.

        # Extract road network from static facts.
        self.road_network = {}
        for fact in static_facts:
            if match(fact, "road", "*", "*"):
                _, l1, l2 = get_parts(fact)
                if l1 not in self.road_network:
                    self.road_network[l1] = set()
                self.road_network[l1].add(l2)

        # Extract capacity constraints from static facts.
        self.capacity_constraints = {}
        for fact in static_facts:
            if match(fact, "capacity-predecessor", "*", "*"):
                _, s1, s2 = get_parts(fact)
                self.capacity_constraints[s2] = s1

        # 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 needed.
            if current_location == goal_location:
                continue

            # Compute the minimal number of actions required to move the package to its goal.
            # Assume that each move between connected locations costs 1 action.
            # If the package is in a vehicle, it may need to be dropped at the goal location.
            # If the package is not in a vehicle, it may need to be picked up by a vehicle and then transported.
            if in_vehicle:
                total_cost += 1  # Drop the package at the goal location.
            else:
                total_cost += 2  # Pick up the package and drop it at the goal location.

            # Add the cost of moving the vehicle to the goal location if necessary.
            # This is a simplified estimate assuming that the vehicle is already at the package's current location.
            total_cost += 1  # Move the vehicle to the goal location.

        return total_cost
