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.
    - Calibrating cameras and taking images of objectives.
    - Communicating data to the lander.

    # Assumptions
    - Each rover can carry only one soil or rock sample at a time.
    - Cameras must be calibrated before taking images.
    - Data can only be communicated to the lander from specific waypoints.

    # Heuristic Initialization
    - Extract goal conditions and static facts from the task.
    - Build data structures to map rovers, waypoints, objectives, and cameras.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current state of each rover:
       - Location (waypoint).
       - Whether it has soil/rock samples.
       - Whether it has images of objectives.
    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 the required image.
         - Estimate the number of actions to navigate to the lander and communicate.
       - If the goal is to collect soil/rock samples:
         - Estimate the number of actions to navigate to the sample location and collect it.
       - If the goal is to take images:
         - Estimate the number of actions to calibrate the camera, navigate to the objective, and take the image.
    3. Sum the estimated actions for all goals to compute the heuristic value.
    """

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

        # Extract static information into suitable data structures.
        self.rovers = set()
        self.waypoints = set()
        self.objectives = set()
        self.cameras = set()
        self.lander = None

        for fact in self.static:
            parts = get_parts(fact)
            if parts[0] == "at_lander":
                self.lander = parts[1]
            elif parts[0] == "rover":
                self.rovers.add(parts[1])
            elif parts[0] == "waypoint":
                self.waypoints.add(parts[1])
            elif parts[0] == "objective":
                self.objectives.add(parts[1])
            elif parts[0] == "camera":
                self.cameras.add(parts[1])

        # Map rovers to their stores.
        self.rover_stores = {}
        for fact in self.static:
            if match(fact, "store_of", "*", "*"):
                store, rover = get_parts(fact)[1:]
                self.rover_stores[rover] = store

        # Map cameras to their calibration targets.
        self.calibration_targets = {}
        for fact in self.static:
            if match(fact, "calibration_target", "*", "*"):
                camera, objective = get_parts(fact)[1:]
                self.calibration_targets[camera] = objective

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

        total_cost = 0

        # Check each goal and compute the cost to achieve it.
        for goal in self.goals:
            parts = get_parts(goal)
            if parts[0] == "communicated_soil_data":
                waypoint = parts[1]
                # Find a rover that has the soil sample.
                for rover in self.rovers:
                    if f"(have_soil_analysis {rover} {waypoint})" in state:
                        # Navigate to the lander and communicate.
                        total_cost += 2  # navigate + communicate
                        break
            elif parts[0] == "communicated_rock_data":
                waypoint = parts[1]
                # Find a rover that has the rock sample.
                for rover in self.rovers:
                    if f"(have_rock_analysis {rover} {waypoint})" in state:
                        # Navigate to the lander and communicate.
                        total_cost += 2  # navigate + communicate
                        break
            elif parts[0] == "communicated_image_data":
                objective, mode = parts[1], parts[2]
                # Find a rover that has the image.
                for rover in self.rovers:
                    if f"(have_image {rover} {objective} {mode})" in state:
                        # Navigate to the lander and communicate.
                        total_cost += 2  # navigate + communicate
                        break
            elif parts[0] == "at_soil_sample":
                waypoint = parts[1]
                # Find a rover that can collect the sample.
                for rover in self.rovers:
                    if f"(equipped_for_soil_analysis {rover})" in state:
                        # Navigate to the waypoint and collect the sample.
                        total_cost += 2  # navigate + sample
                        break
            elif parts[0] == "at_rock_sample":
                waypoint = parts[1]
                # Find a rover that can collect the sample.
                for rover in self.rovers:
                    if f"(equipped_for_rock_analysis {rover})" in state:
                        # Navigate to the waypoint and collect the sample.
                        total_cost += 2  # navigate + sample
                        break
            elif parts[0] == "have_image":
                objective, mode = parts[1], parts[2]
                # Find a rover with the required camera.
                for rover in self.rovers:
                    for camera in self.cameras:
                        if f"(on_board {camera} {rover})" in state and f"(supports {camera} {mode})" in state:
                            # Calibrate, navigate, and take the image.
                            total_cost += 3  # calibrate + navigate + take_image
                            break

        return total_cost
