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.
    - 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 specific waypoints.

    # 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 is carrying soil or rock samples.
       - Whether its camera is calibrated.
    2. Identify the current state of each objective:
       - Whether the required images have been taken.
       - Whether the data has been communicated to the lander.
    3. Identify the current state of soil and rock samples:
       - Whether they have been collected.
       - Whether the data has been communicated to the lander.
    4. For each unachieved goal:
       - Estimate the number of actions required to navigate to the relevant waypoint.
       - Estimate the number of actions required to collect samples or take images.
       - Estimate the number of actions required to communicate data to the lander.
    5. Sum the estimated actions for all unachieved 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 information about waypoints, rovers, cameras, and objectives from static facts.
        self.waypoints = set()
        self.rovers = set()
        self.cameras = set()
        self.objectives = set()
        self.lander_location = None

        for fact in self.static:
            predicate, *args = get_parts(fact)
            if predicate == "at_lander":
                self.lander_location = args[1]
            elif 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])

        # Store goal locations for each objective and sample.
        self.goal_locations = {}
        for goal in self.goals:
            predicate, *args = get_parts(goal)
            if predicate == "communicated_soil_data":
                self.goal_locations[args[0]] = "soil"
            elif predicate == "communicated_rock_data":
                self.goal_locations[args[0]] = "rock"
            elif predicate == "communicated_image_data":
                self.goal_locations[args[0]] = "image"

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

        # Track the current state of rovers, objectives, and samples.
        rover_locations = {}
        rover_samples = {}
        rover_calibrated = {}
        objective_images = {}
        sample_collected = {}

        for fact in state:
            predicate, *args = get_parts(fact)
            if predicate == "at":
                rover, waypoint = args
                rover_locations[rover] = waypoint
            elif predicate == "have_soil_analysis":
                rover, waypoint = args
                rover_samples[rover] = ("soil", waypoint)
            elif predicate == "have_rock_analysis":
                rover, waypoint = args
                rover_samples[rover] = ("rock", waypoint)
            elif predicate == "calibrated":
                camera, rover = args
                rover_calibrated[rover] = True
            elif predicate == "have_image":
                rover, objective, mode = args
                objective_images[(objective, mode)] = True
            elif predicate == "communicated_soil_data":
                waypoint = args[0]
                sample_collected[waypoint] = True
            elif predicate == "communicated_rock_data":
                waypoint = args[0]
                sample_collected[waypoint] = True
            elif predicate == "communicated_image_data":
                objective, mode = args
                objective_images[(objective, mode)] = True

        total_cost = 0

        # Estimate actions for soil and rock samples.
        for waypoint, sample_type in self.goal_locations.items():
            if sample_type in ["soil", "rock"]:
                if waypoint not in sample_collected:
                    # Find a rover that can collect the sample.
                    for rover in self.rovers:
                        if rover_locations.get(rover) == waypoint:
                            # Rover is already at the waypoint.
                            total_cost += 1  # Sample action.
                            break
                    else:
                        # Rover needs to navigate to the waypoint.
                        total_cost += 2  # Navigate and sample actions.

        # Estimate actions for image objectives.
        for (objective, mode), communicated in objective_images.items():
            if not communicated:
                # Find a rover with a calibrated camera.
                for rover in self.rovers:
                    if rover_calibrated.get(rover, False):
                        # Rover can take the image.
                        total_cost += 1  # Take image action.
                        break
                else:
                    # Rover needs to calibrate the camera.
                    total_cost += 2  # Calibrate and take image actions.

        # Estimate actions for communicating data to the lander.
        for waypoint, sample_type in self.goal_locations.items():
            if sample_type in ["soil", "rock"] and waypoint not in sample_collected:
                # Rover needs to communicate the sample data.
                total_cost += 1  # Communicate action.
            elif sample_type == "image":
                # Rover needs to communicate the image data.
                total_cost += 1  # Communicate action.

        return total_cost
