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 in the Miconic domain. It considers the current state of the elevator, the passengers' origins and destinations, and whether passengers have been boarded or served.

    # Assumptions:
    - The elevator 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 and the `destin` relationships for passengers.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current floor of the elevator.
    2. For each passenger:
       - If the passenger is not yet served:
         - If the passenger is not boarded, estimate the cost to board them (move to their origin floor and board).
         - Estimate the cost to serve them (move to their destination floor and depart).
    3. Sum the estimated costs for all passengers to get the total heuristic value.
    """

    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.

        # Map passengers to their destination floors.
        self.destin = {
            get_parts(fact)[1]: get_parts(fact)[2]
            for fact in static_facts
            if match(fact, "destin", "*", "*")
        }

        # Map passengers to their origin floors.
        self.origin = {
            get_parts(fact)[1]: get_parts(fact)[2]
            for fact in static_facts
            if match(fact, "origin", "*", "*")
        }

        # 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", "*", "*")
        }

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

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

        total_cost = 0  # Initialize the heuristic cost.

        for passenger in self.destin:
            # Check if the passenger has already been served.
            if f"(served {passenger})" in state:
                continue

            # If the passenger is not boarded, estimate the cost to board them.
            if f"(boarded {passenger})" not in state:
                origin_floor = self.origin[passenger]
                # Estimate the cost to move to the origin floor.
                if lift_at != origin_floor:
                    total_cost += 1  # Assume one action to move to the origin floor.
                total_cost += 1  # Board the passenger.

            # Estimate the cost to serve the passenger.
            destin_floor = self.destin[passenger]
            if lift_at != destin_floor:
                total_cost += 1  # Assume one action to move to the destination floor.
            total_cost += 1  # Depart the passenger.

        return total_cost
