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:
    - Navigating to waypoints to collect soil and rock samples.
    - Calibrating cameras and taking images of objectives.
    - Communicating data to the lander.

    # Assumptions
    - Each rover can carry only one soil or rock sample at a time.
    - Cameras must be calibrated before taking images.
    - Data can only be communicated to the lander from a visible waypoint.

    # Heuristic Initialization
    - Extract goal conditions and static facts from the task.
    - Identify the locations of the lander, 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 current state of each rover:
       - Location of the rover.
       - Whether it is carrying soil or rock samples.
       - Whether its camera is calibrated.
       - Whether it has taken required images.
    2. For each goal condition:
       - If the goal is to communicate soil data, estimate the number of actions needed to collect and communicate the soil sample.
       - If the goal is to communicate rock data, estimate the number of actions needed to collect and communicate the rock sample.
       - If the goal is to communicate image data, estimate the number of actions needed to calibrate the camera, take the image, and communicate it.
    3. Sum the estimated actions for all goal conditions 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 location
        self.lander_location = None
        for fact in self.static:
            if match(fact, "at_lander", "*", "*"):
                parts = get_parts(fact)
                self.lander_location = parts[2]

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

        # 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 calibration targets
        self.calibration_targets = {}
        for fact in self.static:
            if match(fact, "calibration_target", "*", "*"):
                parts = get_parts(fact)
                camera = parts[1]
                objective = parts[2]
                self.calibration_targets[camera] = objective

        # Extract camera supports
        self.camera_supports = {}
        for fact in self.static:
            if match(fact, "supports", "*", "*"):
                parts = get_parts(fact)
                camera = parts[1]
                mode = parts[2]
                if camera not in self.camera_supports:
                    self.camera_supports[camera] = set()
                self.camera_supports[camera].add(mode)

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

        total_cost = 0

        # Check each goal condition
        for goal in self.goals:
            parts = get_parts(goal)
            if parts[0] == "communicated_soil_data":
                waypoint = parts[1]
                # Estimate actions to collect and communicate soil data
                total_cost += self.estimate_soil_actions(state, waypoint)
            elif parts[0] == "communicated_rock_data":
                waypoint = parts[1]
                # Estimate actions to collect and communicate rock data
                total_cost += self.estimate_rock_actions(state, waypoint)
            elif parts[0] == "communicated_image_data":
                objective = parts[1]
                mode = parts[2]
                # Estimate actions to calibrate, take image, and communicate
                total_cost += self.estimate_image_actions(state, objective, mode)

        return total_cost

    def estimate_soil_actions(self, state, waypoint):
        """Estimate the number of actions needed to collect and communicate soil data."""
        cost = 0

        # Check if soil data is already communicated
        if f"(communicated_soil_data {waypoint})" in state:
            return 0

        # Find a rover that can collect soil
        for rover in self.rover_capabilities:
            if "soil_analysis" in self.rover_capabilities[rover]:
                # Check if the rover is at the waypoint
                if f"(at {rover} {waypoint})" in state:
                    # Collect soil sample
                    cost += 1
                    # Communicate soil data
                    cost += 1
                    break

        return cost

    def estimate_rock_actions(self, state, waypoint):
        """Estimate the number of actions needed to collect and communicate rock data."""
        cost = 0

        # Check if rock data is already communicated
        if f"(communicated_rock_data {waypoint})" in state:
            return 0

        # Find a rover that can collect rock
        for rover in self.rover_capabilities:
            if "rock_analysis" in self.rover_capabilities[rover]:
                # Check if the rover is at the waypoint
                if f"(at {rover} {waypoint})" in state:
                    # Collect rock sample
                    cost += 1
                    # Communicate rock data
                    cost += 1
                    break

        return cost

    def estimate_image_actions(self, state, objective, mode):
        """Estimate the number of actions needed to calibrate, take image, and communicate."""
        cost = 0

        # Check if image data is already communicated
        if f"(communicated_image_data {objective} {mode})" in state:
            return 0

        # Find a rover with a camera that supports the mode
        for rover in self.rover_capabilities:
            if "imaging" in self.rover_capabilities[rover]:
                # Find a camera on the rover that supports the mode
                for camera in self.camera_supports:
                    if mode in self.camera_supports[camera] and f"(on_board {camera} {rover})" in state:
                        # Calibrate the camera
                        cost += 1
                        # Take the image
                        cost += 1
                        # Communicate the image data
                        cost += 1
                        break

        return cost
