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 communication requires the rover to be at a waypoint visible to the lander.

    # 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 a waypoint visible to the lander.
       - If the goal is to communicate image data:
         - Check if the rover has the required image.
         - Estimate the number of actions to navigate to a waypoint visible to the lander.
       - 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 an image:
         - Estimate the number of actions to calibrate the camera, navigate to the objective's visible waypoint, 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_location = None
        self.visible_from = {}  # Maps objectives to waypoints from which they are visible.
        self.calibration_targets = {}  # Maps cameras to their calibration targets.
        self.supports = {}  # Maps cameras to supported modes.

        for fact in self.static:
            predicate, *args = get_parts(fact)
            if predicate == "at_lander":
                self.lander_location = args[1]
            elif predicate == "visible_from":
                objective, waypoint = args
                if objective not in self.visible_from:
                    self.visible_from[objective] = set()
                self.visible_from[objective].add(waypoint)
            elif predicate == "calibration_target":
                camera, objective = args
                self.calibration_targets[camera] = objective
            elif predicate == "supports":
                camera, mode = args
                if camera not in self.supports:
                    self.supports[camera] = set()
                self.supports[camera].add(mode)
            elif predicate == "on_board":
                camera, rover = args
                self.cameras.add(camera)
                self.rovers.add(rover)
            elif predicate == "store_of":
                store, rover = args
                self.rovers.add(rover)
            elif predicate == "can_traverse":
                rover, waypoint1, waypoint2 = args
                self.rovers.add(rover)
                self.waypoints.add(waypoint1)
                self.waypoints.add(waypoint2)

    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:
            predicate, *args = get_parts(goal)
            if predicate == "communicated_soil_data":
                waypoint = args[0]
                # Find a rover that has the soil sample and is at a waypoint visible to the lander.
                for rover in self.rovers:
                    if f"(have_soil_analysis {rover} {waypoint})" in state:
                        # Check if the rover is at a waypoint visible to the lander.
                        rover_location = None
                        for fact in state:
                            if match(fact, "at", rover, "*"):
                                rover_location = get_parts(fact)[2]
                                break
                        if rover_location and f"(visible {rover_location} {self.lander_location})" in self.static:
                            total_cost += 1  # Communicate soil data.
                        else:
                            total_cost += 2  # Navigate to a visible waypoint and communicate.
            elif predicate == "communicated_rock_data":
                waypoint = args[0]
                # Similar logic as for soil data.
                for rover in self.rovers:
                    if f"(have_rock_analysis {rover} {waypoint})" in state:
                        rover_location = None
                        for fact in state:
                            if match(fact, "at", rover, "*"):
                                rover_location = get_parts(fact)[2]
                                break
                        if rover_location and f"(visible {rover_location} {self.lander_location})" in self.static:
                            total_cost += 1
                        else:
                            total_cost += 2
            elif predicate == "communicated_image_data":
                objective, mode = args
                # Find a rover that has the image and is at a waypoint visible to the lander.
                for rover in self.rovers:
                    if f"(have_image {rover} {objective} {mode})" in state:
                        rover_location = None
                        for fact in state:
                            if match(fact, "at", rover, "*"):
                                rover_location = get_parts(fact)[2]
                                break
                        if rover_location and f"(visible {rover_location} {self.lander_location})" in self.static:
                            total_cost += 1
                        else:
                            total_cost += 2
            elif predicate == "have_soil_analysis":
                rover, waypoint = args
                # Estimate the cost to navigate to the waypoint and collect the sample.
                total_cost += 2  # Navigate and sample.
            elif predicate == "have_rock_analysis":
                rover, waypoint = args
                # Similar logic as for soil analysis.
                total_cost += 2
            elif predicate == "have_image":
                rover, objective, mode = args
                # Estimate the cost to calibrate the camera, navigate to the objective's visible waypoint, and take the image.
                total_cost += 3  # Calibrate, navigate, and take image.

        return total_cost
