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., "(origin p1 f1)".
    - `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 MiconicHeuristic(Heuristic):
    """
    A domain-dependent heuristic for the Miconic domain.

    # Summary
    This heuristic estimates the number of actions required to serve all passengers by considering:
    - The current position of the lift.
    - The origin and destination floors of each passenger.
    - Whether passengers are already boarded or served.

    # Assumptions
    - The lift can move between floors using the `up` and `down` actions.
    - Passengers must be boarded at their origin floor and served at their destination floor.
    - The heuristic does not need to be admissible, so it can overestimate the number of actions.

    # Heuristic Initialization
    - Extract the goal conditions (all passengers must be served).
    - Extract static facts, such as the `above` relationships between floors.
    - Create a mapping of passengers to their origin and destination floors.

    # Step-By-Step Thinking for Computing Heuristic
    1. For each passenger:
       - If the passenger is already served, no actions are needed.
       - If the passenger is boarded, the lift must move to their destination floor and perform a `depart` action.
       - If the passenger is not boarded, the lift must move to their origin floor, perform a `board` action, then move to their destination floor and perform a `depart` action.
    2. Calculate the number of floor changes required for the lift to reach the necessary floors.
    3. Sum the actions required for all passengers to estimate the total cost.
    """

    def __init__(self, task):
        """Initialize the heuristic by extracting goal conditions and static facts."""
        self.goals = task.goals  # Goal conditions.
        static_facts = task.static  # Facts that are not affected by actions.

        # Extract passenger origin and destination floors from static facts.
        self.origin_floors = {}
        self.destin_floors = {}
        for fact in static_facts:
            parts = get_parts(fact)
            if parts[0] == "origin":
                passenger, floor = parts[1], parts[2]
                self.origin_floors[passenger] = floor
            elif parts[0] == "destin":
                passenger, floor = parts[1], parts[2]
                self.destin_floors[passenger] = floor

        # Extract the `above` relationships between floors.
        self.above = set()
        for fact in static_facts:
            parts = get_parts(fact)
            if parts[0] == "above":
                floor1, floor2 = parts[1], parts[2]
                self.above.add((floor1, floor2))

    def __call__(self, node):
        """Estimate the number of actions required to serve all passengers."""
        state = node.state  # Current world state.

        # Track which passengers are boarded or served.
        boarded_passengers = set()
        served_passengers = set()
        for fact in state:
            parts = get_parts(fact)
            if parts[0] == "boarded":
                boarded_passengers.add(parts[1])
            elif parts[0] == "served":
                served_passengers.add(parts[1])

        # Get the current floor of the lift.
        lift_at = None
        for fact in state:
            parts = get_parts(fact)
            if parts[0] == "lift-at":
                lift_at = parts[1]
                break

        total_cost = 0  # Initialize action cost counter.

        for passenger in self.origin_floors:
            if passenger in served_passengers:
                continue  # Passenger is already served.

            origin_floor = self.origin_floors[passenger]
            destin_floor = self.destin_floors[passenger]

            if passenger in boarded_passengers:
                # Passenger is already boarded; need to move to destination and depart.
                if lift_at != destin_floor:
                    # Calculate the number of floor changes required.
                    total_cost += self._calculate_floor_changes(lift_at, destin_floor)
                total_cost += 1  # Depart action.
            else:
                # Passenger is not boarded; need to move to origin, board, then move to destination and depart.
                if lift_at != origin_floor:
                    total_cost += self._calculate_floor_changes(lift_at, origin_floor)
                total_cost += 1  # Board action.
                if origin_floor != destin_floor:
                    total_cost += self._calculate_floor_changes(origin_floor, destin_floor)
                total_cost += 1  # Depart action.

        return total_cost

    def _calculate_floor_changes(self, floor1, floor2):
        """
        Calculate the number of floor changes required to move from `floor1` to `floor2`.
        """
        if floor1 == floor2:
            return 0
        # Assume each floor change requires one action (either `up` or `down`).
        return 1  # Simplified for this heuristic; could be improved with a more accurate floor distance calculation.
