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

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

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current state of each rover:
       - Location (waypoint).
       - Whether it has soil or rock samples.
       - Whether it has images of objectives.
    2. Identify the current state of each objective:
       - Whether its image data has been communicated.
    3. Identify the current state of each waypoint:
       - Whether soil or rock samples are available.
    4. Compute the cost for each rover to achieve its goals:
       - Navigate to waypoints to collect samples.
       - Calibrate cameras and take images.
       - Communicate data to the lander.
    5. Sum the costs 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 waypoints, rovers, cameras, and objectives from static facts.
        self.waypoints = set()
        self.rovers = set()
        self.cameras = set()
        self.objectives = set()

        for fact in self.static:
            predicate, *args = get_parts(fact)
            if predicate == "waypoint":
                self.waypoints.add(args[0])
            elif predicate == "rover":
                self.rovers.add(args[0])
            elif predicate == "camera":
                self.cameras.add(args[0])
            elif predicate == "objective":
                self.objectives.add(args[0])

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

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

        # Map objectives to their visible waypoints.
        self.visible_from = {}
        for fact in self.static:
            if match(fact, "visible_from", "*", "*"):
                objective, waypoint = get_parts(fact)[1:]
                if objective not in self.visible_from:
                    self.visible_from[objective] = set()
                self.visible_from[objective].add(waypoint)

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

        # Initialize cost.
        total_cost = 0

        # Check for each goal whether it has been achieved.
        for goal in self.goals:
            if goal not in state:
                predicate, *args = get_parts(goal)
                if predicate == "communicated_soil_data":
                    waypoint = args[0]
                    total_cost += self._cost_to_communicate_soil(waypoint, state)
                elif predicate == "communicated_rock_data":
                    waypoint = args[0]
                    total_cost += self._cost_to_communicate_rock(waypoint, state)
                elif predicate == "communicated_image_data":
                    objective, mode = args
                    total_cost += self._cost_to_communicate_image(objective, mode, state)

        return total_cost

    def _cost_to_communicate_soil(self, waypoint, state):
        """Compute the cost to communicate soil data from a waypoint."""
        cost = 0

        # Find a rover that can collect and communicate the soil sample.
        for rover in self.rovers:
            if f"(have_soil_analysis {rover} {waypoint})" in state:
                # Soil sample already collected, just need to communicate.
                cost += 1  # Communicate action.
                return cost

        # Soil sample not collected yet.
        cost += 1  # Navigate to waypoint.
        cost += 1  # Sample soil.
        cost += 1  # Communicate soil data.
        return cost

    def _cost_to_communicate_rock(self, waypoint, state):
        """Compute the cost to communicate rock data from a waypoint."""
        cost = 0

        # Find a rover that can collect and communicate the rock sample.
        for rover in self.rovers:
            if f"(have_rock_analysis {rover} {waypoint})" in state:
                # Rock sample already collected, just need to communicate.
                cost += 1  # Communicate action.
                return cost

        # Rock sample not collected yet.
        cost += 1  # Navigate to waypoint.
        cost += 1  # Sample rock.
        cost += 1  # Communicate rock data.
        return cost

    def _cost_to_communicate_image(self, objective, mode, state):
        """Compute the cost to communicate image data for an objective."""
        cost = 0

        # Check if the image has already been taken.
        for rover in self.rovers:
            if f"(have_image {rover} {objective} {mode})" in state:
                # Image already taken, just need to communicate.
                cost += 1  # Communicate action.
                return cost

        # Image not taken yet.
        cost += 1  # Navigate to a waypoint visible from the objective.
        cost += 1  # Calibrate camera.
        cost += 1  # Take image.
        cost += 1  # Communicate image data.
        return cost
