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 following tasks:
    - Navigating rovers to waypoints with soil/rock samples or objectives.
    - Collecting soil/rock samples.
    - Taking images of objectives.
    - Communicating data to the lander.

    # Assumptions
    - Rovers can only carry one sample at a time (soil or rock).
    - 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 to their stores, cameras, and waypoints.
    - Identify the lander's location and the visibility graph between waypoints.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current state of each rover:
       - Location.
       - Whether it has a soil/rock sample.
       - Whether it has taken required images.
    2. For each goal condition:
       - If the goal is to communicate soil/rock data:
         - Check if the rover has the 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 image.
         - Estimate the number of actions to navigate to a waypoint visible to the lander.
    3. Sum the estimated actions for all goal conditions.
    """

    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

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

        # 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):
        """Estimate the number of actions required to achieve the goal state."""
        state = node.state
        total_cost = 0

        # Check each goal condition
        for goal in self.goals:
            predicate, *args = get_parts(goal)
            if predicate == "communicated_soil_data":
                waypoint = args[0]
                # Find a rover with the soil sample
                for fact in state:
                    if match(fact, "have_soil_analysis", "*", waypoint):
                        rover = get_parts(fact)[1]
                        # Estimate navigation cost to a waypoint visible to the lander
                        total_cost += self.estimate_navigation_cost(rover, state)
                        break
            elif predicate == "communicated_rock_data":
                waypoint = args[0]
                # Find a rover with the rock sample
                for fact in state:
                    if match(fact, "have_rock_analysis", "*", waypoint):
                        rover = get_parts(fact)[1]
                        # Estimate navigation cost to a waypoint visible to the lander
                        total_cost += self.estimate_navigation_cost(rover, state)
                        break
            elif predicate == "communicated_image_data":
                objective, mode = args[0], args[1]
                # Find a rover with the image
                for fact in state:
                    if match(fact, "have_image", "*", objective, mode):
                        rover = get_parts(fact)[1]
                        # Estimate navigation cost to a waypoint visible to the lander
                        total_cost += self.estimate_navigation_cost(rover, state)
                        break

        return total_cost

    def estimate_navigation_cost(self, rover, state):
        """Estimate the number of navigation actions required for a rover to reach a waypoint visible to the lander."""
        # Find the rover's current location
        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

        # Check if the rover is already at a waypoint visible to the lander
        if rover_location in self.visibility and self.lander_location in self.visibility[rover_location]:
            return 0

        # Estimate the number of navigation steps required
        # This is a simple heuristic: assume 1 action per waypoint in the shortest path
        # In practice, a more accurate estimate could use a graph search algorithm
        return 1
