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 navigation required to reach waypoints for sampling, imaging, and communication.

    # Assumptions
    - Each rover can carry only one sample at a time (soil or rock).
    - Each rover can take images only if its camera is calibrated.
    - Communication with the lander requires the rover to be at a waypoint visible from the lander's location.

    # Heuristic Initialization
    - Extract goal conditions and static facts from the task.
    - Build data structures to map rovers to their stores, cameras, and waypoints.
    - Identify the lander's location and the visibility graph for waypoints.

    # 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 rover:
        - Calculate the navigation cost to reach the required waypoints for sampling, imaging, and communication.
        - Add the cost of sampling, imaging, and communication actions.
    4. Sum the costs for all rovers to get the total 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 the lander's location.
        self.lander_location = None
        for fact in self.static:
            if match(fact, "at_lander", "*", "*"):
                self.lander_location = get_parts(fact)[2]
                break

        # Build a visibility graph for waypoints.
        self.visible = {}
        for fact in self.static:
            if match(fact, "visible", "*", "*"):
                parts = get_parts(fact)
                waypoint1, waypoint2 = parts[1], parts[2]
                if waypoint1 not in self.visible:
                    self.visible[waypoint1] = set()
                self.visible[waypoint1].add(waypoint2)

        # Map rovers to their stores and cameras.
        self.rover_stores = {}
        self.rover_cameras = {}
        for fact in self.static:
            if match(fact, "store_of", "*", "*"):
                store, rover = get_parts(fact)[1], get_parts(fact)[2]
                self.rover_stores[rover] = store
            elif match(fact, "on_board", "*", "*"):
                camera, rover = get_parts(fact)[1], get_parts(fact)[2]
                if rover not in self.rover_cameras:
                    self.rover_cameras[rover] = set()
                self.rover_cameras[rover].add(camera)

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

        # Initialize the total cost.
        total_cost = 0

        # Check for soil and rock samples that need to be communicated.
        for goal in self.goals:
            if match(goal, "communicated_soil_data", "*"):
                waypoint = get_parts(goal)[1]
                if not any(match(fact, "communicated_soil_data", waypoint) for fact in state):
                    total_cost += 2  # Sample and communicate.
            elif match(goal, "communicated_rock_data", "*"):
                waypoint = get_parts(goal)[1]
                if not any(match(fact, "communicated_rock_data", waypoint) for fact in state):
                    total_cost += 2  # Sample and communicate.

        # Check for images that need to be communicated.
        for goal in self.goals:
            if match(goal, "communicated_image_data", "*", "*"):
                objective, mode = get_parts(goal)[1], get_parts(goal)[2]
                if not any(match(fact, "communicated_image_data", objective, mode) for fact in state):
                    total_cost += 3  # Calibrate, take image, and communicate.

        # Add navigation costs for rovers.
        for rover in self.rover_stores:
            # Check if the rover needs to move to a waypoint for sampling or imaging.
            if any(match(fact, "at", rover, "*") for fact in state):
                current_waypoint = next(get_parts(fact)[2] for fact in state if match(fact, "at", rover, "*"))
                # Estimate the cost to reach the lander's location for communication.
                if current_waypoint != self.lander_location:
                    total_cost += 1  # Navigate to the lander's location.

        return total_cost
