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 sample at a time (soil or rock).
    - Cameras must be calibrated before taking images.
    - Communication with the lander requires the rover to be at a visible waypoint.

    # Heuristic Initialization
    - Extract goal conditions and static facts from the task.
    - Map waypoints to their visible waypoints and objectives.
    - Identify which rovers are equipped for soil/rock analysis and imaging.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current state of each rover:
       - Location.
       - Whether it has soil/rock samples or images.
       - Whether its camera is 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 the lander and communicate.
       - If the goal is to communicate image data:
         - Check if the rover has the required image.
         - Estimate the number of actions to navigate to the lander and communicate.
       - If the goal is to collect a soil/rock sample:
         - Estimate the number of actions to navigate to the sample location and collect it.
       - If the goal is to take an image:
         - Estimate the number of actions to calibrate the camera, navigate to the objective, 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
        self.static = task.static

        # Extract visible waypoints and objectives from static facts.
        self.visible = {}
        self.visible_from = {}
        for fact in self.static:
            parts = get_parts(fact)
            if parts[0] == "visible":
                self.visible.setdefault(parts[1], set()).add(parts[2])
            elif parts[0] == "visible_from":
                self.visible_from.setdefault(parts[1], set()).add(parts[2])

        # Extract rover capabilities from static facts.
        self.equipped_for_soil = set()
        self.equipped_for_rock = set()
        self.equipped_for_imaging = set()
        for fact in self.static:
            parts = get_parts(fact)
            if parts[0] == "equipped_for_soil_analysis":
                self.equipped_for_soil.add(parts[1])
            elif parts[0] == "equipped_for_rock_analysis":
                self.equipped_for_rock.add(parts[1])
            elif parts[0] == "equipped_for_imaging":
                self.equipped_for_imaging.add(parts[1])

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

        # Check each goal and compute the required actions.
        for goal in self.goals:
            parts = get_parts(goal)
            if parts[0] == "communicated_soil_data":
                waypoint = parts[1]
                # Find a rover with the soil sample.
                for fact in state:
                    if match(fact, "have_soil_analysis", "*", waypoint):
                        rover = get_parts(fact)[1]
                        # Estimate actions to navigate to the lander and communicate.
                        total_cost += self._estimate_communication_cost(rover, waypoint, state)
                        break
            elif parts[0] == "communicated_rock_data":
                waypoint = parts[1]
                # Find a rover with the rock sample.
                for fact in state:
                    if match(fact, "have_rock_analysis", "*", waypoint):
                        rover = get_parts(fact)[1]
                        # Estimate actions to navigate to the lander and communicate.
                        total_cost += self._estimate_communication_cost(rover, waypoint, state)
                        break
            elif parts[0] == "communicated_image_data":
                objective, mode = parts[1], parts[2]
                # Find a rover with the image.
                for fact in state:
                    if match(fact, "have_image", "*", objective, mode):
                        rover = get_parts(fact)[1]
                        # Estimate actions to navigate to the lander and communicate.
                        total_cost += self._estimate_communication_cost(rover, None, state)
                        break
            elif parts[0] == "at_soil_sample":
                waypoint = parts[1]
                # Estimate actions to navigate to the waypoint and collect the sample.
                total_cost += self._estimate_sample_collection_cost(waypoint, state)
            elif parts[0] == "at_rock_sample":
                waypoint = parts[1]
                # Estimate actions to navigate to the waypoint and collect the sample.
                total_cost += self._estimate_sample_collection_cost(waypoint, state)
            elif parts[0] == "have_image":
                objective, mode = parts[1], parts[2]
                # Estimate actions to calibrate the camera, navigate to the objective, and take the image.
                total_cost += self._estimate_imaging_cost(objective, mode, state)

        return total_cost

    def _estimate_communication_cost(self, rover, waypoint, state):
        """Estimate the number of actions to communicate data to the lander."""
        cost = 0
        # Find the lander's location.
        lander_location = None
        for fact in state:
            if match(fact, "at_lander", "*", "*"):
                lander_location = get_parts(fact)[2]
                break
        if not lander_location:
            return float('inf')  # No lander found.

        # 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 not rover_location:
            return float('inf')  # Rover not found.

        # Estimate navigation cost.
        if rover_location != lander_location:
            cost += 1  # Navigate to the lander.

        # Communication action.
        cost += 1
        return cost

    def _estimate_sample_collection_cost(self, waypoint, state):
        """Estimate the number of actions to collect a soil/rock sample."""
        cost = 0
        # Find a rover capable of collecting the sample.
        rover = None
        for fact in state:
            if match(fact, "at", "*", waypoint):
                rover = get_parts(fact)[1]
                break
        if not rover:
            return float('inf')  # No rover at the waypoint.

        # Check if the rover is equipped for the task.
        if waypoint in self.visible_from.get("at_soil_sample", set()):
            if rover not in self.equipped_for_soil:
                return float('inf')  # Rover not equipped for soil analysis.
        elif waypoint in self.visible_from.get("at_rock_sample", set()):
            if rover not in self.equipped_for_rock:
                return float('inf')  # Rover not equipped for rock analysis.

        # Sample collection action.
        cost += 1
        return cost

    def _estimate_imaging_cost(self, objective, mode, state):
        """Estimate the number of actions to take an image."""
        cost = 0
        # Find a rover with the required camera.
        rover = None
        camera = None
        for fact in state:
            if match(fact, "on_board", "*", "*"):
                camera_part, rover_part = get_parts(fact)[1], get_parts(fact)[2]
                if rover_part in self.equipped_for_imaging:
                    rover = rover_part
                    camera = camera_part
                    break
        if not rover or not camera:
            return float('inf')  # No suitable rover or camera found.

        # Check if the camera is calibrated.
        calibrated = False
        for fact in state:
            if match(fact, "calibrated", camera, rover):
                calibrated = True
                break
        if not calibrated:
            cost += 1  # Calibrate the camera.

        # Navigate to the objective's visible waypoint.
        cost += 1

        # Take the image.
        cost += 1
        return cost
