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., "(at box1 loc_1_1)".
    - `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 SokobanHeuristic(Heuristic):
    """
    A domain-dependent heuristic for the Sokoban domain.

    # Summary
    This heuristic estimates the number of actions required to push all boxes to their goal locations.
    It considers the distance between the robot and the boxes, as well as the distance between the boxes and their goal locations.

    # Assumptions
    - The robot can only push one box at a time.
    - The robot must be adjacent to a box to push it.
    - The robot can move freely to any clear location.
    - The heuristic does not account for deadlocks or blocked paths.

    # Heuristic Initialization
    - Extract goal locations for each box from the task goals.
    - Extract adjacency relationships between locations from the static facts.
    - Compute the shortest path distances between all pairs of locations using a breadth-first search (BFS).

    # Step-By-Step Thinking for Computing Heuristic
    1. For each box, compute the distance from its current location to its goal location.
    2. For each box, compute the distance from the robot's current location to the box's current location.
    3. Sum the distances for all boxes, weighted by the number of actions required to move the robot and push the boxes.
    """

    def __init__(self, task):
        """
        Initialize the heuristic by extracting:
        - Goal locations for each box.
        - Adjacency relationships between locations.
        - Compute shortest path distances between all pairs of locations.
        """
        self.goals = task.goals  # Goal conditions.
        static_facts = task.static  # Facts that are not affected by actions.

        # Extract goal locations for each box.
        self.goal_locations = {}
        for goal in self.goals:
            predicate, *args = get_parts(goal)
            if predicate == "at":
                box, location = args
                self.goal_locations[box] = location

        # Build adjacency graph from static facts.
        self.adjacency = {}
        for fact in static_facts:
            predicate, *args = get_parts(fact)
            if predicate == "adjacent":
                loc1, loc2, _ = args
                if loc1 not in self.adjacency:
                    self.adjacency[loc1] = set()
                if loc2 not in self.adjacency:
                    self.adjacency[loc2] = set()
                self.adjacency[loc1].add(loc2)
                self.adjacency[loc2].add(loc1)

        # Compute shortest path distances between all pairs of locations using BFS.
        self.distances = {}
        for loc in self.adjacency:
            self.distances[loc] = self._bfs(loc)

    def _bfs(self, start):
        """Compute shortest path distances from `start` to all other locations using BFS."""
        distances = {start: 0}
        queue = [start]
        while queue:
            current = queue.pop(0)
            for neighbor in self.adjacency[current]:
                if neighbor not in distances:
                    distances[neighbor] = distances[current] + 1
                    queue.append(neighbor)
        return distances

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

        # Extract the robot's current location.
        robot_location = None
        for fact in state:
            if match(fact, "at-robot", "*"):
                robot_location = get_parts(fact)[1]
                break

        # Extract the current locations of all boxes.
        box_locations = {}
        for fact in state:
            if match(fact, "at", "*", "*"):
                box, location = get_parts(fact)[1:]
                box_locations[box] = location

        total_cost = 0  # Initialize action cost counter.

        for box, goal_location in self.goal_locations.items():
            current_location = box_locations[box]

            # Distance from the robot to the box.
            robot_to_box = self.distances[robot_location][current_location]

            # Distance from the box to its goal.
            box_to_goal = self.distances[current_location][goal_location]

            # Total cost for this box: robot moves to the box, then pushes it to the goal.
            total_cost += robot_to_box + box_to_goal

        return total_cost
