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 before they can be served.
    - 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 and the `destin` of each passenger.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current state of the lift and the passengers:
       - Check which passengers are already served (no further actions needed).
       - Check which passengers are boarded (they need to be served at their destination).
       - Check which passengers are still at their origin (they need to be boarded and served).
    2. For each passenger not yet served:
       - Calculate the number of floors the lift needs to travel to reach their origin (if not boarded) or destination (if boarded).
       - Add the cost of boarding or serving the passenger.
    3. Sum the total number of actions required to serve all passengers.
    """

    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 the `above` relationships between floors.
        self.above = {
            (get_parts(fact)[1], get_parts(fact)[2])
            for fact in static_facts
            if match(fact, "above", "*", "*")
        }

        # Extract the destination floors for each passenger.
        self.destin = {
            get_parts(fact)[1]: get_parts(fact)[2]
            for fact in static_facts
            if match(fact, "destin", "*", "*")
        }

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

        # Identify which passengers are already served.
        served_passengers = {
            get_parts(fact)[1]
            for fact in state
            if match(fact, "served", "*")
        }

        # Identify which passengers are boarded.
        boarded_passengers = {
            get_parts(fact)[1]
            for fact in state
            if match(fact, "boarded", "*")
        }

        # Identify the current floor of the lift.
        lift_at = next(
            get_parts(fact)[1]
            for fact in state
            if match(fact, "lift-at", "*")
        )

        total_cost = 0  # Initialize the heuristic cost.

        for passenger, destin_floor in self.destin.items():
            if passenger in served_passengers:
                continue  # No further actions needed for this passenger.

            # Determine the current floor of the passenger.
            if passenger in boarded_passengers:
                # Passenger is already boarded; need to go to destination.
                current_floor = lift_at
                target_floor = destin_floor
            else:
                # Passenger is at origin; need to go to origin first.
                origin_fact = next(
                    fact for fact in state if match(fact, "origin", passenger, "*")
                )
                current_floor = get_parts(origin_fact)[2]
                target_floor = current_floor

            # Calculate the number of floors to travel.
            if current_floor != target_floor:
                # Count the number of floors between current and target.
                # Since the `above` relationship is transitive, we can assume a direct path.
                total_cost += abs(int(current_floor[1:]) - int(target_floor[1:]))

            # Add the cost of boarding or serving the passenger.
            if passenger not in boarded_passengers:
                total_cost += 1  # Boarding action.
            total_cost += 1  # Serving action.

        return total_cost
