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 rover1 waypoint1)".
    - `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 RoversHeuristic(Heuristic):
    """
    A domain-dependent heuristic for the Rovers domain.

    # Summary
    This heuristic estimates the number of actions required to achieve all goals in the Rovers domain.
    It considers the following tasks:
    - Navigating to waypoints to collect soil/rock samples.
    - Taking images of objectives.
    - Communicating data to the lander.

    # Assumptions
    - Each rover can carry only one sample at a time (soil or rock).
    - Each rover can take images of objectives if equipped with the appropriate camera.
    - Communication with the lander requires the rover to be at a waypoint visible to the lander.

    # Heuristic Initialization
    - Extract goal conditions and static facts from the task.
    - Identify the lander's location, waypoints, and objectives.
    - Map rovers to their stores, cameras, and capabilities.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current state of each rover:
       - Location of the rover.
       - Whether it has soil/rock samples.
       - Whether it has taken required images.
    2. For each goal:
       - If the goal is to communicate soil/rock data:
         - Check if the rover has the required sample.
         - Estimate the number of actions to navigate to the lander and communicate.
       - If the goal is to communicate image data:
         - Check if the rover has taken the required image.
         - Estimate the number of actions to navigate to the lander and communicate.
    3. Sum the estimated actions for all goals.
    """

    def __init__(self, task):
        """Initialize the heuristic by extracting goal conditions and static facts."""
        self.goals = task.goals  # Goal conditions.
        static_facts = task.static  # Static facts.

        # Extract lander location.
        self.lander_location = None
        for fact in static_facts:
            if match(fact, "at_lander", "*", "*"):
                _, lander, location = get_parts(fact)
                self.lander_location = location
                break

        # Map rovers to their stores and capabilities.
        self.rovers = {}
        for fact in static_facts:
            if match(fact, "store_of", "*", "*"):
                _, store, rover = get_parts(fact)
                if rover not in self.rovers:
                    self.rovers[rover] = {"store": store, "capabilities": set()}
            elif match(fact, "equipped_for_*", "*"):
                _, capability, rover = get_parts(fact)
                if rover in self.rovers:
                    self.rovers[rover]["capabilities"].add(capability)

        # Map cameras to their rovers and objectives.
        self.cameras = {}
        for fact in static_facts:
            if match(fact, "on_board", "*", "*"):
                _, camera, rover = get_parts(fact)
                self.cameras[camera] = {"rover": rover, "objectives": set()}
            elif match(fact, "calibration_target", "*", "*"):
                _, camera, objective = get_parts(fact)
                if camera in self.cameras:
                    self.cameras[camera]["objectives"].add(objective)

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

        # Check if all goals are already satisfied.
        if self.goals <= state:
            return 0

        total_cost = 0  # Initialize action cost counter.

        # Iterate over all goals.
        for goal in self.goals:
            predicate, *args = get_parts(goal)

            if predicate == "communicated_soil_data":
                waypoint = args[0]
                # Find a rover with the required soil sample.
                for rover in self.rovers:
                    if f"(have_soil_analysis {rover} {waypoint})" in state:
                        # Estimate actions to communicate the data.
                        total_cost += self._estimate_communication_cost(rover, state)
                        break

            elif predicate == "communicated_rock_data":
                waypoint = args[0]
                # Find a rover with the required rock sample.
                for rover in self.rovers:
                    if f"(have_rock_analysis {rover} {waypoint})" in state:
                        # Estimate actions to communicate the data.
                        total_cost += self._estimate_communication_cost(rover, state)
                        break

            elif predicate == "communicated_image_data":
                objective, mode = args
                # Find a rover with the required image.
                for rover in self.rovers:
                    if f"(have_image {rover} {objective} {mode})" in state:
                        # Estimate actions to communicate the data.
                        total_cost += self._estimate_communication_cost(rover, state)
                        break

        return total_cost

    def _estimate_communication_cost(self, rover, state):
        """
        Estimate the number of actions required for a rover to communicate data to the lander.

        Args:
            rover: The rover that needs to communicate.
            state: The current state of the world.

        Returns:
            int: The estimated number of actions.
        """
        # Find the rover's current location.
        rover_location = None
        for fact in state:
            if match(fact, "at", rover, "*"):
                _, _, location = get_parts(fact)
                rover_location = location
                break

        if not rover_location:
            return 0  # Rover is not at any waypoint.

        # Check if the rover is already at a waypoint visible to the lander.
        if f"(visible {rover_location} {self.lander_location})" in state:
            return 1  # Only the communication action is needed.

        # Otherwise, estimate the number of navigation actions required.
        return 2  # Navigate to a visible waypoint and communicate.
