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.
    - Identify the lander's location, rover capabilities, and waypoint visibility.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current state of each rover:
       - Location of each rover.
       - Whether the rover has soil or rock samples.
       - Whether the rover 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 and is at a waypoint visible to the lander.
       - If the goal is to communicate rock data, check if the rover has the rock sample and is at a waypoint visible to the lander.
       - If the goal is to communicate image data, check if the rover has taken the required image and is at a waypoint visible to the lander.
    3. Estimate the number of actions required to achieve each goal:
       - Navigate to the required waypoint.
       - Collect samples or take images.
       - Communicate data to the lander.
    4. Sum the estimated actions for all goals to get 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 rover capabilities
        self.rover_capabilities = {}
        for fact in self.static:
            if match(fact, "equipped_for_*", "*"):
                parts = get_parts(fact)
                rover = parts[1]
                capability = parts[0].split("_")[2]
                if rover not in self.rover_capabilities:
                    self.rover_capabilities[rover] = set()
                self.rover_capabilities[rover].add(capability)

        # Extract waypoint visibility
        self.visible_waypoints = {}
        for fact in self.static:
            if match(fact, "visible", "*", "*"):
                parts = get_parts(fact)
                from_waypoint = parts[1]
                to_waypoint = parts[2]
                if from_waypoint not in self.visible_waypoints:
                    self.visible_waypoints[from_waypoint] = set()
                self.visible_waypoints[from_waypoint].add(to_waypoint)

    def __call__(self, node):
        """Estimate the number of actions required to achieve the goal state."""
        state = node.state

        # Initialize heuristic value
        heuristic_value = 0

        # Check each goal condition
        for goal in self.goals:
            if match(goal, "communicated_soil_data", "*"):
                waypoint = get_parts(goal)[1]
                # Check if any rover has the soil sample and is at a waypoint visible to the lander
                for fact in state:
                    if match(fact, "have_soil_analysis", "*", waypoint):
                        rover = get_parts(fact)[1]
                        # Check if the rover is at a waypoint visible to the lander
                        for loc_fact in state:
                            if match(loc_fact, "at", rover, "*"):
                                rover_location = get_parts(loc_fact)[2]
                                if rover_location in self.visible_waypoints.get(self.lander_location, set()):
                                    break
                        else:
                            heuristic_value += 1  # Need to navigate to a visible waypoint
                        break
                else:
                    heuristic_value += 2  # Need to collect soil and navigate

            elif match(goal, "communicated_rock_data", "*"):
                waypoint = get_parts(goal)[1]
                # Check if any rover has the rock sample and is at a waypoint visible to the lander
                for fact in state:
                    if match(fact, "have_rock_analysis", "*", waypoint):
                        rover = get_parts(fact)[1]
                        # Check if the rover is at a waypoint visible to the lander
                        for loc_fact in state:
                            if match(loc_fact, "at", rover, "*"):
                                rover_location = get_parts(loc_fact)[2]
                                if rover_location in self.visible_waypoints.get(self.lander_location, set()):
                                    break
                        else:
                            heuristic_value += 1  # Need to navigate to a visible waypoint
                        break
                else:
                    heuristic_value += 2  # Need to collect rock and navigate

            elif match(goal, "communicated_image_data", "*", "*"):
                objective = get_parts(goal)[1]
                mode = get_parts(goal)[2]
                # Check if any rover has the image and is at a waypoint visible to the lander
                for fact in state:
                    if match(fact, "have_image", "*", objective, mode):
                        rover = get_parts(fact)[1]
                        # Check if the rover is at a waypoint visible to the lander
                        for loc_fact in state:
                            if match(loc_fact, "at", rover, "*"):
                                rover_location = get_parts(loc_fact)[2]
                                if rover_location in self.visible_waypoints.get(self.lander_location, set()):
                                    break
                        else:
                            heuristic_value += 1  # Need to navigate to a visible waypoint
                        break
                else:
                    heuristic_value += 2  # Need to take image and navigate

        return heuristic_value
