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 the
    relationships between floors (e.g., which floors are above others).

    # 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 assumes that the elevator can only carry one passenger at a time (simplified for efficiency).

    # Heuristic Initialization
    - Extract the goal conditions (all passengers must be served).
    - Extract static facts, such as the `above` relationships between floors and the `destin` (destination) of each passenger.
    - Build a mapping of passengers to their destinations and a graph of floor relationships.

    # 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:
           - Add the cost of moving the elevator to the passenger's origin floor.
           - Add the cost of boarding the passenger.
         - Add the cost of moving the elevator to the passenger's destination floor.
         - Add the cost of serving the passenger.
    3. Sum the costs for all passengers to get the total heuristic value.
    """

    def __init__(self, task):
        """
        Initialize the heuristic by extracting:
        - Goal conditions (all passengers must be served).
        - Static facts (`above` relationships and passenger destinations).
        """
        self.goals = task.goals  # Goal conditions.
        static_facts = task.static  # Facts that are not affected by actions.

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

        # Build a graph of floor relationships using "above" facts.
        self.above_graph = {}
        for fact in static_facts:
            if match(fact, "above", "*", "*"):
                floor1, floor2 = get_parts(fact)[1], get_parts(fact)[2]
                if floor1 not in self.above_graph:
                    self.above_graph[floor1] = set()
                self.above_graph[floor1].add(floor2)

    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.
        elevator_floor = None
        for fact in state:
            if match(fact, "lift-at", "*"):
                elevator_floor = get_parts(fact)[1]
                break
        assert elevator_floor is not None, "Elevator floor not found in state."

        total_cost = 0  # Initialize action cost counter.

        for passenger, destination in self.passenger_destinations.items():
            # Check if the passenger is already served.
            if f"(served {passenger})" in state:
                continue

            # Check if the passenger is boarded.
            if f"(boarded {passenger})" not in state:
                # Passenger is not boarded; need to move to origin floor and board.
                origin_floor = None
                for fact in state:
                    if match(fact, "origin", passenger, "*"):
                        origin_floor = get_parts(fact)[2]
                        break
                assert origin_floor is not None, f"Origin floor not found for passenger {passenger}."

                # Add cost of moving to origin floor.
                total_cost += self._compute_floor_distance(elevator_floor, origin_floor)
                # Add cost of boarding.
                total_cost += 1

            # Add cost of moving to destination floor.
            total_cost += self._compute_floor_distance(elevator_floor, destination)
            # Add cost of serving.
            total_cost += 1

        return total_cost

    def _compute_floor_distance(self, floor1, floor2):
        """
        Compute the minimum number of elevator moves required to travel from `floor1` to `floor2`.

        - `floor1`: The starting floor.
        - `floor2`: The target floor.
        - Returns the number of `up` or `down` actions required.
        """
        if floor1 == floor2:
            return 0

        # Perform a breadth-first search to find the shortest path.
        visited = set()
        queue = [(floor1, 0)]
        while queue:
            current_floor, distance = queue.pop(0)
            if current_floor == floor2:
                return distance
            if current_floor in visited:
                continue
            visited.add(current_floor)

            # Explore floors above and below.
            if current_floor in self.above_graph:
                for neighbor in self.above_graph[current_floor]:
                    queue.append((neighbor, distance + 1))
            # Reverse lookup for floors below.
            for floor, neighbors in self.above_graph.items():
                if current_floor in neighbors:
                    queue.append((floor, distance + 1))

        raise ValueError(f"No path found from {floor1} to {floor2}.")
