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/rock samples or take images.
    - Communicating data to the lander.
    - Calibrating cameras for imaging.

    # Assumptions
    - Each rover can carry only one soil or rock sample at a time.
    - 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, waypoints, cameras, and objectives.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current state of each rover:
       - Location (waypoint).
       - Whether it has soil/rock samples or images.
       - Whether its cameras are calibrated.
    2. For each goal:
       - If the goal is to communicate soil/rock data:
         - Check if the rover has the required 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 required image.
         - Estimate the number of actions to navigate to a waypoint visible to the lander.
       - If the goal is to collect soil/rock samples:
         - Estimate the number of actions to navigate to the sample location and collect it.
       - If the goal is to take images:
         - Estimate the number of actions to calibrate the camera, navigate to the objective's location, and take the image.
    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  # Goal conditions.
        static_facts = task.static  # Facts that are not affected by actions.

        # Extract static information into suitable data structures.
        self.visible = {
            (get_parts(fact)[1], get_parts(fact)[2]
            for fact in static_facts
            if match(fact, "visible", "*", "*")
        }
        self.can_traverse = {
            (get_parts(fact)[1], get_parts(fact)[2], get_parts(fact)[3])
            for fact in static_facts
            if match(fact, "can_traverse", "*", "*", "*")
        }
        self.visible_from = {
            (get_parts(fact)[1], get_parts(fact)[2])
            for fact in static_facts
            if match(fact, "visible_from", "*", "*")
        }
        self.calibration_targets = {
            (get_parts(fact)[1], get_parts(fact)[2]
            for fact in static_facts
            if match(fact, "calibration_target", "*", "*")
        }
        self.supports = {
            (get_parts(fact)[1], get_parts(fact)[2]
            for fact in static_facts
            if match(fact, "supports", "*", "*")
        }

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

        # If the state is a goal state, return 0.
        if self.goals <= state:
            return 0

        total_cost = 0  # Initialize action cost counter.

        # Check for soil/rock data communication goals.
        for goal in self.goals:
            if match(goal, "communicated_soil_data", "*"):
                waypoint = get_parts(goal)[1]
                # Check if any rover has the required soil sample.
                has_sample = any(
                    match(fact, "have_soil_analysis", "*", waypoint) for fact in state
                )
                if not has_sample:
                    # Estimate the cost to collect the sample.
                    total_cost += 2  # Navigate to the waypoint and collect the sample.
                # Estimate the cost to communicate the data.
                total_cost += 1  # Navigate to a waypoint visible to the lander.

            elif match(goal, "communicated_rock_data", "*"):
                waypoint = get_parts(goal)[1]
                # Check if any rover has the required rock sample.
                has_sample = any(
                    match(fact, "have_rock_analysis", "*", waypoint) for fact in state
                )
                if not has_sample:
                    # Estimate the cost to collect the sample.
                    total_cost += 2  # Navigate to the waypoint and collect the sample.
                # Estimate the cost to communicate the data.
                total_cost += 1  # Navigate to a waypoint visible to the lander.

            elif match(goal, "communicated_image_data", "*", "*"):
                objective, mode = get_parts(goal)[1], get_parts(goal)[2]
                # Check if any rover has the required image.
                has_image = any(
                    match(fact, "have_image", "*", objective, mode) for fact in state
                )
                if not has_image:
                    # Estimate the cost to take the image.
                    total_cost += 3  # Calibrate, navigate, and take the image.
                # Estimate the cost to communicate the data.
                total_cost += 1  # Navigate to a waypoint visible to the lander.

        return total_cost
