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 before they can be served.
    - The elevator must be at the correct floor to board or serve a passenger.

    # 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.
    - Create data structures to store the relationships between floors and passengers' destinations.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current floor of the elevator.
    2. For each passenger:
       - If the passenger has already been served, no actions are needed.
       - If the passenger has been boarded but not served, the elevator must move to the destination floor and perform a `depart` action.
       - If the passenger has not been boarded, the elevator must move to the origin floor, perform a `board` action, then move to the destination floor and perform a `depart` action.
    3. Sum the number of actions required for 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 `destin` relationships for passengers.
        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 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 action cost counter.

        # Iterate over all passengers.
        for fact in state:
            if match(fact, "origin", "*", "*"):
                passenger, origin_floor = get_parts(fact)[1], get_parts(fact)[2]
                destin_floor = self.destin[passenger]

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

                # Check if the passenger has been boarded.
                if f"(boarded {passenger})" in state:
                    # The elevator must move to the destination floor and perform a `depart` action.
                    if lift_at != destin_floor:
                        total_cost += 1  # Move to the destination floor.
                    total_cost += 1  # Perform the `depart` action.
                else:
                    # The elevator must move to the origin floor, perform a `board` action, then move to the destination floor and perform a `depart` action.
                    if lift_at != origin_floor:
                        total_cost += 1  # Move to the origin floor.
                    total_cost += 1  # Perform the `board` action.
                    if origin_floor != destin_floor:
                        total_cost += 1  # Move to the destination floor.
                    total_cost += 1  # Perform the `depart` action.

        return total_cost
