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.

    The heuristic is computed by summing the estimated actions for each goal condition, considering the current state of the rovers.

    # Assumptions
    - Each rover can carry only one sample at a time (soil or rock).
    - Each rover can take images of multiple objectives, but only one at a time.
    - 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. **Soil and Rock Samples**:
       - For each required soil or rock sample, check if it has already been collected and communicated.
       - If not, estimate the number of actions needed to collect and communicate the sample:
         - Navigate to the sample location.
         - Collect the sample.
         - Navigate to a waypoint visible to the lander.
         - Communicate the sample data.

    2. **Image Data**:
       - For each required image, check if it has already been taken and communicated.
       - If not, estimate the number of actions needed to take and communicate the image:
         - Calibrate the camera (if not already calibrated).
         - Navigate to a waypoint visible to the objective.
         - Take the image.
         - Navigate to a waypoint visible to the lander.
         - Communicate the image data.

    3. **Summing Actions**:
       - The total heuristic value is the sum of all estimated actions for soil, rock, and image tasks.
    """

    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.

        # Map rovers to their stores.
        self.rover_stores = {
            get_parts(fact)[1]: get_parts(fact)[2]
            for fact in static_facts
            if match(fact, "store_of", "*", "*")
        }

        # Map cameras to their calibration targets.
        self.camera_targets = {
            get_parts(fact)[1]: get_parts(fact)[2]
            for fact in static_facts
            if match(fact, "calibration_target", "*", "*")
        }

        # Map rovers to their onboard cameras.
        self.rover_cameras = {
            get_parts(fact)[1]: get_parts(fact)[2]
            for fact in static_facts
            if match(fact, "on_board", "*", "*")
        }

        # Identify the lander's location.
        self.lander_location = next(
            get_parts(fact)[2]
            for fact in static_facts
            if match(fact, "at_lander", "*", "*")
        )

        # Build a visibility graph between waypoints.
        self.visibility_graph = {
            (get_parts(fact)[1], get_parts(fact)[2])
            for fact in static_facts
            if match(fact, "visible", "*", "*")
        }

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

        total_cost = 0  # Initialize the heuristic cost.

        # Check soil and rock sample goals.
        for goal in self.goals:
            if match(goal, "communicated_soil_data", "*"):
                waypoint = get_parts(goal)[1]
                if goal not in state:
                    # Estimate actions to collect and communicate soil sample.
                    total_cost += self._estimate_soil_sample_actions(waypoint, state)
            elif match(goal, "communicated_rock_data", "*"):
                waypoint = get_parts(goal)[1]
                if goal not in state:
                    # Estimate actions to collect and communicate rock sample.
                    total_cost += self._estimate_rock_sample_actions(waypoint, state)
            elif match(goal, "communicated_image_data", "*", "*"):
                objective, mode = get_parts(goal)[1], get_parts(goal)[2]
                if goal not in state:
                    # Estimate actions to take and communicate image.
                    total_cost += self._estimate_image_actions(objective, mode, state)

        return total_cost

    def _estimate_soil_sample_actions(self, waypoint, state):
        """Estimate actions to collect and communicate a soil sample."""
        cost = 0

        # Find a rover equipped for soil analysis.
        rover = next(
            get_parts(fact)[1]
            for fact in state
            if match(fact, "equipped_for_soil_analysis", "*")
        )

        # Navigate to the sample location.
        cost += self._estimate_navigation_cost(rover, waypoint, state)

        # Collect the sample.
        cost += 1  # sample_soil action.

        # Navigate to a waypoint visible to the lander.
        cost += self._estimate_navigation_cost(rover, self.lander_location, state)

        # Communicate the sample data.
        cost += 1  # communicate_soil_data action.

        return cost

    def _estimate_rock_sample_actions(self, waypoint, state):
        """Estimate actions to collect and communicate a rock sample."""
        cost = 0

        # Find a rover equipped for rock analysis.
        rover = next(
            get_parts(fact)[1]
            for fact in state
            if match(fact, "equipped_for_rock_analysis", "*")
        )

        # Navigate to the sample location.
        cost += self._estimate_navigation_cost(rover, waypoint, state)

        # Collect the sample.
        cost += 1  # sample_rock action.

        # Navigate to a waypoint visible to the lander.
        cost += self._estimate_navigation_cost(rover, self.lander_location, state)

        # Communicate the sample data.
        cost += 1  # communicate_rock_data action.

        return cost

    def _estimate_image_actions(self, objective, mode, state):
        """Estimate actions to take and communicate an image."""
        cost = 0

        # Find a rover equipped for imaging.
        rover = next(
            get_parts(fact)[1]
            for fact in state
            if match(fact, "equipped_for_imaging", "*")
        )

        # Find the camera on the rover.
        camera = self.rover_cameras[rover]

        # Calibrate the camera (if not already calibrated).
        if f"(calibrated {camera} {rover})" not in state:
            cost += 1  # calibrate action.

        # Navigate to a waypoint visible to the objective.
        waypoint = next(
            get_parts(fact)[2]
            for fact in state
            if match(fact, "visible_from", objective, "*")
        )
        cost += self._estimate_navigation_cost(rover, waypoint, state)

        # Take the image.
        cost += 1  # take_image action.

        # Navigate to a waypoint visible to the lander.
        cost += self._estimate_navigation_cost(rover, self.lander_location, state)

        # Communicate the image data.
        cost += 1  # communicate_image_data action.

        return cost

    def _estimate_navigation_cost(self, rover, target_waypoint, state):
        """Estimate the number of navigate actions required to reach a waypoint."""
        # Find the current location of the rover.
        current_waypoint = next(
            get_parts(fact)[2]
            for fact in state
            if match(fact, "at", rover, "*")
        )

        # If already at the target, no cost.
        if current_waypoint == target_waypoint:
            return 0

        # Otherwise, assume one navigate action is sufficient.
        return 1
