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, and packages can be picked up or dropped based on capacity constraints.

    # Heuristic Initialization
    - Extract goal locations for each package from the task goals.
    - Extract static information about roads and capacity relationships from the task's 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. Calculate the minimal number of actions required to move the package to its goal:
       - If the package is inside a vehicle, it must be dropped at the goal location.
       - If the package is at a location, it must be picked up by a vehicle and then dropped at the goal location.
    4. Consider the capacity constraints of vehicles when calculating the number of pick-up and drop actions.
    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.
        - Static facts (roads and capacity relationships).
        """
        self.goals = task.goals  # Goal conditions.
        static_facts = task.static  # Facts that are not affected by actions.

        # Extract roads from static facts.
        self.roads = set()
        for fact in static_facts:
            if match(fact, "road", "*", "*"):
                parts = get_parts(fact)
                self.roads.add((parts[1], parts[2]))

        # Extract capacity relationships from static facts.
        self.capacity_predecessors = {}
        for fact in static_facts:
            if match(fact, "capacity-predecessor", "*", "*"):
                parts = get_parts(fact)
                self.capacity_predecessors[parts[1]] = parts[2]

        # 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:
                # If the package is inside a vehicle, it must be dropped at the goal location.
                total_cost += 1  # Drop action.
            else:
                # If the package is at a location, it must be picked up and then dropped.
                total_cost += 1  # Pick-up action.
                total_cost += 1  # Drop action.

            # Calculate the minimal number of drive actions required to move the vehicle to the goal location.
            if in_vehicle:
                vehicle_location = current_locations[current_location]
                if vehicle_location != goal_location:
                    # Find the shortest path between the current location and the goal location.
                    # For simplicity, assume each road connection requires one drive action.
                    total_cost += 1  # Drive action.

        return total_cost
