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

    # Assumptions
    - Each rover can carry only one sample at a time (soil or rock).
    - Each rover can take images of objectives if equipped with the appropriate camera.
    - Communication with the lander 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 of the rover.
       - Whether it has soil or rock samples.
       - Whether it has taken required images.
    2. For each goal condition:
       - 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 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", "*", "*"):
                waypoint1, waypoint2 = get_parts(fact)[1], get_parts(fact)[2]
                if waypoint1 not in self.visibility:
                    self.visibility[waypoint1] = set()
                self.visibility[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
            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):
        """Estimate the number of actions required to achieve the goal state."""
        state = node.state
        total_cost = 0

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

        # Iterate over each goal condition
        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:
                        # Estimate actions to communicate
                        total_cost += self.estimate_communication_cost(rover, state)
            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:
                        # Estimate actions to communicate
                        total_cost += self.estimate_communication_cost(rover, state)
            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:
                        # Estimate actions to communicate
                        total_cost += self.estimate_communication_cost(rover, state)

        return total_cost

    def estimate_communication_cost(self, rover, state):
        """Estimate the number of actions required for a rover to communicate data to the lander."""
        cost = 0

        # 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 is None:
            return float('inf')  # Rover is not at any waypoint

        # 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]:
            cost += 1  # Communicate action
        else:
            # Estimate navigation actions to reach a visible waypoint
            cost += 2  # Navigate to a visible waypoint and communicate

        return cost
