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 waypoints, rovers, cameras, and objectives.

    # 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 or images.
       - Whether its camera is calibrated.
    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 the data.
       - 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 the image.
       - If the goal is to collect a soil/rock sample:
         - 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, 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
        self.waypoints = set()
        self.rovers = set()
        self.cameras = set()
        self.objectives = set()
        self.lander_location = None
        self.visible_from = {}  # Maps objectives to waypoints
        self.calibration_targets = {}  # Maps cameras to objectives
        self.supports = {}  # Maps cameras to supported modes

        for fact in self.static:
            parts = get_parts(fact)
            if parts[0] == "visible_from":
                objective, waypoint = parts[1], parts[2]
                if objective not in self.visible_from:
                    self.visible_from[objective] = set()
                self.visible_from[objective].add(waypoint)
            elif parts[0] == "calibration_target":
                camera, objective = parts[1], parts[2]
                self.calibration_targets[camera] = objective
            elif parts[0] == "supports":
                camera, mode = parts[1], parts[2]
                if camera not in self.supports:
                    self.supports[camera] = set()
                self.supports[camera].add(mode)
            elif parts[0] == "at_lander":
                self.lander_location = parts[2]

        # Extract objects from the task
        for obj in task.facts:
            parts = get_parts(obj)
            if parts[0] == "waypoint":
                self.waypoints.add(parts[1])
            elif parts[0] == "rover":
                self.rovers.add(parts[1])
            elif parts[0] == "camera":
                self.cameras.add(parts[1])
            elif parts[0] == "objective":
                self.objectives.add(parts[1])

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

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

        # Extract current state information
        rover_locations = {}
        rover_samples = {}
        rover_images = {}
        calibrated_cameras = set()

        for fact in state:
            parts = get_parts(fact)
            if parts[0] == "at":
                rover, waypoint = parts[1], parts[2]
                rover_locations[rover] = waypoint
            elif parts[0] == "have_soil_analysis":
                rover, waypoint = parts[1], parts[2]
                rover_samples[rover] = ("soil", waypoint)
            elif parts[0] == "have_rock_analysis":
                rover, waypoint = parts[1], parts[2]
                rover_samples[rover] = ("rock", waypoint)
            elif parts[0] == "have_image":
                rover, objective, mode = parts[1], parts[2], parts[3]
                if rover not in rover_images:
                    rover_images[rover] = set()
                rover_images[rover].add((objective, mode))
            elif parts[0] == "calibrated":
                camera, rover = parts[1], parts[2]
                calibrated_cameras.add((camera, rover))

        # Compute cost for each goal
        for goal in self.goals:
            parts = get_parts(goal)
            if parts[0] == "communicated_soil_data":
                waypoint = parts[1]
                # Find a rover with the soil sample
                for rover, sample in rover_samples.items():
                    if sample[0] == "soil" and sample[1] == waypoint:
                        # Navigate to lander and communicate
                        total_cost += 2  # navigate + communicate
                        break
            elif parts[0] == "communicated_rock_data":
                waypoint = parts[1]
                # Find a rover with the rock sample
                for rover, sample in rover_samples.items():
                    if sample[0] == "rock" and sample[1] == waypoint:
                        # Navigate to lander and communicate
                        total_cost += 2  # navigate + communicate
                        break
            elif parts[0] == "communicated_image_data":
                objective, mode = parts[1], parts[2]
                # Find a rover with the image
                for rover, images in rover_images.items():
                    if (objective, mode) in images:
                        # Navigate to lander and communicate
                        total_cost += 2  # navigate + communicate
                        break
            elif parts[0] == "at_soil_sample":
                waypoint = parts[1]
                # Navigate to waypoint and sample soil
                total_cost += 2  # navigate + sample
            elif parts[0] == "at_rock_sample":
                waypoint = parts[1]
                # Navigate to waypoint and sample rock
                total_cost += 2  # navigate + sample
            elif parts[0] == "have_image":
                objective, mode = parts[1], parts[2]
                # Calibrate camera, navigate to objective, and take image
                total_cost += 3  # calibrate + navigate + take_image

        return total_cost
