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 still need to be collected and communicated.
    - The number of images that still need to be taken and communicated.
    - The distance rovers need to travel to achieve these tasks.

    # Assumptions
    - Rovers can only carry one sample at a time (soil or rock).
    - Rovers must return to a lander to communicate data.
    - The heuristic assumes that rovers can navigate directly to the required waypoints without obstacles.

    # Heuristic Initialization
    - Extract goal conditions and static facts from the task.
    - Identify the locations of landers, waypoints, and objectives.
    - Determine which rovers are equipped for soil analysis, rock analysis, and imaging.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the number of soil and rock samples that still need to be collected and communicated.
    2. Identify the number of images that still need to be taken and communicated.
    3. For each rover, calculate the distance to the nearest waypoint where a sample or image can be collected.
    4. Estimate the number of actions required to navigate to the waypoint, collect the sample or image, and communicate the data.
    5. Sum the estimated actions 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 lander locations
        self.lander_locations = {
            get_parts(fact)[1]: get_parts(fact)[2]
            for fact in self.static
            if match(fact, "at_lander", "*", "*")
        }

        # Extract rover capabilities
        self.rover_capabilities = {
            "soil_analysis": set(),
            "rock_analysis": set(),
            "imaging": set(),
        }
        for fact in self.static:
            if match(fact, "equipped_for_soil_analysis", "*"):
                self.rover_capabilities["soil_analysis"].add(get_parts(fact)[1])
            elif match(fact, "equipped_for_rock_analysis", "*"):
                self.rover_capabilities["rock_analysis"].add(get_parts(fact)[1])
            elif match(fact, "equipped_for_imaging", "*"):
                self.rover_capabilities["imaging"].add(get_parts(fact)[1])

        # Extract waypoint visibility
        self.visible_waypoints = {
            (get_parts(fact)[1], get_parts(fact)[2])
            for fact in self.static
            if match(fact, "visible", "*", "*")
        }

        # Extract calibration targets
        self.calibration_targets = {
            get_parts(fact)[1]: get_parts(fact)[2]
            for fact in self.static
            if match(fact, "calibration_target", "*", "*")
        }

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

        # Initialize the heuristic value
        total_cost = 0

        # Check for soil data communication goals
        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  # Navigate to waypoint and communicate

        # Check for rock data communication goals
        for goal in self.goals:
            if 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  # Navigate to waypoint and communicate

        # Check for image data communication goals
        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

        return total_cost
