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 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 of the rover.
       - Whether it has soil or rock samples.
       - Whether it has images of objectives.
    2. Identify the current state of each objective:
       - Whether the required images have been taken and communicated.
    3. Identify the current state of each waypoint:
       - Whether soil or rock samples are still available.
    4. Compute the cost for each rover to achieve its goals:
       - Navigate to waypoints to collect samples.
       - Calibrate cameras and take images.
       - Navigate to a waypoint visible to the lander to communicate data.
    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 information from static facts
        self.waypoints = set()
        self.rovers = set()
        self.cameras = set()
        self.objectives = set()
        self.lander = None

        for fact in self.static:
            parts = get_parts(fact)
            if parts[0] == "waypoint":
                self.waypoints.add(parts[1])
            elif parts[0] == "rover":
                self.rovers.add(parts[1])
            elif parts[0] == "camera":
                self.cameras.add(parts[1])
            elif parts[0] == "objective":
                self.objectives.add(parts[1])
            elif parts[0] == "lander":
                self.lander = parts[1]

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

        # Build a map from cameras to their supported modes
        self.supports = {}
        for fact in self.static:
            if match(fact, "supports", "*", "*"):
                parts = get_parts(fact)
                camera, mode = parts[1], parts[2]
                if camera not in self.supports:
                    self.supports[camera] = set()
                self.supports[camera].add(mode)

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

        # Initialize the total cost
        total_cost = 0

        # Check if all goals are already satisfied
        if self.goals <= state:
            return 0

        # Iterate over each rover and compute its contribution to the heuristic
        for rover in self.rovers:
            # Find the current location of the rover
            rover_location = None
            for fact in state:
                if match(fact, "at", rover, "*"):
                    rover_location = get_parts(fact)[2]
                    break

            # Check if the rover has any soil or rock samples
            has_soil = any(match(fact, "have_soil_analysis", rover, "*") for fact in state)
            has_rock = any(match(fact, "have_rock_analysis", rover, "*") for fact in state)

            # Check if the rover has any images
            has_images = {}
            for objective in self.objectives:
                for mode in self.supports.values():
                    if any(match(fact, "have_image", rover, objective, mode) for fact in state):
                        has_images[objective] = True
                        break

            # Compute the cost for the rover to achieve its goals
            if has_soil or has_rock:
                # If the rover has samples, it needs to communicate them
                total_cost += 1  # Navigate to a waypoint visible to the lander
                total_cost += 1  # Communicate the data

            for objective in self.objectives:
                if not has_images.get(objective, False):
                    # If the rover doesn't have the required image, it needs to take it
                    total_cost += 1  # Calibrate the camera
                    total_cost += 1  # Take the image
                    total_cost += 1  # Navigate to a waypoint visible to the lander
                    total_cost += 1  # Communicate the image data

        return total_cost
