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-robot 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 heuristic does not account for deadlocks or blocked paths.

    # Heuristic Initialization
    - Extract goal locations for each box from the task goals.
    - Extract the adjacency graph from the static facts to compute distances.

    # Step-By-Step Thinking for Computing Heuristic
    1. For each box, compute the Manhattan distance between its current location and its goal location.
    2. For each box, compute the Manhattan distance between the robot's current location and the box's current location.
    3. Sum these distances to estimate the total number of moves required.
    4. If the robot is already adjacent to a box, reduce the estimated cost by 1 (since no additional move is needed to reach the box).
    """

    def __init__(self, task):
        """
        Initialize the heuristic by extracting:
        - Goal locations for each box.
        - Adjacency graph from static facts.
        """
        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)

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

        # Extract current locations of the robot and boxes.
        robot_location = None
        box_locations = {}
        for fact in state:
            predicate, *args = get_parts(fact)
            if predicate == "at-robot":
                robot_location = args[0]
            elif predicate == "at":
                box, location = args
                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]

            # Compute Manhattan distance between box and goal.
            x1, y1 = map(int, current_location.split("_")[1:])
            x2, y2 = map(int, goal_location.split("_")[1:])
            distance_to_goal = abs(x1 - x2) + abs(y1 - y2)

            # Compute Manhattan distance between robot and box.
            xr, yr = map(int, robot_location.split("_")[1:])
            distance_to_box = abs(xr - x1) + abs(yr - y1)

            # Add distances to total cost.
            total_cost += distance_to_goal + distance_to_box

            # If the robot is adjacent to the box, reduce the cost by 1.
            if current_location in self.adjacency.get(robot_location, set()):
                total_cost -= 1

        return total_cost
