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 and rock samples.
    - Taking images of objectives.
    - Communicating data to the lander.

    # Assumptions
    - Each rover can carry only one soil or rock sample at a time.
    - Each rover can take images of objectives only if the camera is calibrated.
    - Communication of data requires the rover to be at a waypoint visible from the lander.

    # 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 relationships between waypoints.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current state of each rover:
       - Location of the rover.
       - Whether it has soil or rock samples.
       - Whether it has taken required images.
       - Whether it has communicated data to the lander.

    2. For each goal:
       - If the goal is to communicate soil data:
         - Check if the rover has the soil sample.
         - Estimate the number of actions to navigate to the lander and communicate.
       - If the goal is to communicate rock data:
         - Check if the rover has the rock 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 taken the required image.
         - Estimate the number of actions to navigate to the lander and communicate.

    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 lander location
        self.lander_location = None
        for fact in self.static:
            if match(fact, "at_lander", "*", "*"):
                self.lander_location = get_parts(fact)[2]
                break

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

        # Extract rover 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
            if 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 cost
        total_cost = 0

        # Check each goal
        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
                for rover in self.rover_stores:
                    if f"(have_soil_analysis {rover} {waypoint})" in state:
                        # Navigate to lander and communicate
                        total_cost += self.estimate_navigation_cost(rover, state)
                        total_cost += 1  # Communicate action
                        break
            elif predicate == "communicated_rock_data":
                waypoint = args[0]
                # Find a rover that has the rock sample
                for rover in self.rover_stores:
                    if f"(have_rock_analysis {rover} {waypoint})" in state:
                        # Navigate to lander and communicate
                        total_cost += self.estimate_navigation_cost(rover, state)
                        total_cost += 1  # Communicate action
                        break
            elif predicate == "communicated_image_data":
                objective, mode = args[0], args[1]
                # Find a rover that has the image
                for rover in self.rover_cameras:
                    if f"(have_image {rover} {objective} {mode})" in state:
                        # Navigate to lander and communicate
                        total_cost += self.estimate_navigation_cost(rover, state)
                        total_cost += 1  # Communicate action
                        break

        return total_cost

    def estimate_navigation_cost(self, rover, state):
        """Estimate the number of actions required to navigate to the lander."""
        # Find the current location of the rover
        rover_location = None
        for fact in state:
            if match(fact, "at", rover, "*"):
                rover_location = get_parts(fact)[2]
                break

        if rover_location == self.lander_location:
            return 0

        # Estimate the number of navigation steps
        # This is a simple heuristic: assume 1 action per waypoint
        # A more accurate heuristic could use a shortest path algorithm
        return 1
