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 the goal state by considering:
    - The number of soil and rock samples that need to be collected and communicated.
    - The number of images that need to be taken and communicated.
    - The distance between the rover's current location and the locations of samples, objectives, and the lander.

    # Assumptions
    - The rover can only carry one sample at a time (soil or rock).
    - The rover must return to the lander to communicate data.
    - The rover must be at the correct waypoint to take images of objectives.

    # 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 number of soil and rock samples that need to be collected and communicated.
    2. Identify the number of images that need to be taken and communicated.
    3. For each sample or image, estimate the number of actions required:
       - Navigate to the sample or objective location.
       - Perform the sample or image action.
       - Navigate to the lander.
       - Communicate the data.
    4. Sum the estimated actions for all samples and images.
    """

    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

        for fact in self.static:
            parts = get_parts(fact)
            if parts[0] == "at_lander":
                self.lander_location = parts[2]
            elif 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):
        """Estimate the number of actions required to reach the goal state."""
        state = node.state

        # Check if the goal is already reached
        if self.goals <= state:
            return 0

        # Initialize cost
        total_cost = 0

        # Count the number of soil and rock samples that need to be communicated
        soil_samples_needed = sum(1 for goal in self.goals if match(goal, "communicated_soil_data", "*"))
        rock_samples_needed = sum(1 for goal in self.goals if match(goal, "communicated_rock_data", "*"))

        # Count the number of images that need to be communicated
        images_needed = sum(1 for goal in self.goals if match(goal, "communicated_image_data", "*", "*"))

        # Estimate actions for soil samples
        total_cost += soil_samples_needed * 4  # navigate, sample, navigate, communicate

        # Estimate actions for rock samples
        total_cost += rock_samples_needed * 4  # navigate, sample, navigate, communicate

        # Estimate actions for images
        total_cost += images_needed * 5  # navigate, calibrate, take_image, navigate, communicate

        return total_cost
