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 are already 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
      between passengers and their destination floors.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current floor of the elevator.
    2. For each passenger:
       - If the passenger is already served, no actions are needed.
       - If the passenger is boarded, the elevator must move to their destination floor and serve them.
       - If the passenger is not boarded, the elevator must move to their origin floor, board them, then move
         to their destination floor and serve them.
    3. Sum the number of actions required for all passengers, considering the elevator's current position.
    4. Return the total estimated number of actions.
    """

    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 destination floors for each passenger.
        self.destin_floors = {
            get_parts(fact)[1]: get_parts(fact)[2]
            for fact in static_facts
            if match(fact, "destin", "*", "*")
        }

        # Extract origin floors for each passenger.
        self.origin_floors = {
            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_relationships = {
            (get_parts(fact)[1], get_parts(fact)[2])
            for fact in static_facts
            if match(fact, "above", "*", "*")
        }

    def __call__(self, node):
        """Compute an estimate of the minimal number of required actions."""
        state = node.state  # Current world state.

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

        total_cost = 0  # Initialize action cost counter.

        for passenger in self.destin_floors.keys():
            # Check if the passenger is already served.
            if f"(served {passenger})" in state:
                continue

            # Get the passenger's origin and destination floors.
            origin_floor = self.origin_floors[passenger]
            destin_floor = self.destin_floors[passenger]

            # Check if the passenger is already boarded.
            if f"(boarded {passenger})" in state:
                # Elevator must move to the destination floor and serve the passenger.
                total_cost += self._compute_movement_cost(current_floor, destin_floor)
                total_cost += 1  # Serve the passenger.
                current_floor = destin_floor
            else:
                # Elevator must move to the origin floor, board the passenger, then move to the destination floor.
                total_cost += self._compute_movement_cost(current_floor, origin_floor)
                total_cost += 1  # Board the passenger.
                total_cost += self._compute_movement_cost(origin_floor, destin_floor)
                total_cost += 1  # Serve the passenger.
                current_floor = destin_floor

        return total_cost

    def _compute_movement_cost(self, from_floor, to_floor):
        """
        Compute the number of actions required to move the elevator from `from_floor` to `to_floor`.

        - `from_floor`: The current floor of the elevator.
        - `to_floor`: The target floor.
        - Returns the number of `up` or `down` actions required.
        """
        if from_floor == to_floor:
            return 0

        # Determine the direction of movement.
        if (from_floor, to_floor) in self.above_relationships:
            return 1  # One `up` action.
        elif (to_floor, from_floor) in self.above_relationships:
            return 1  # One `down` action.
        else:
            # If floors are not directly connected, assume the worst-case scenario.
            return abs(int(from_floor[1:]) - int(to_floor[1:]))
