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 number of soil and rock samples that need to be collected and communicated.
    - The number of images that need to be taken and communicated.
    - The distance rovers need to travel to achieve these tasks.

    # Assumptions
    - Each rover can carry only one soil or rock sample at a time.
    - Each rover can take images only if its camera is calibrated.
    - Communication of data requires the rover to be at a waypoint visible from 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 number of soil and rock samples that need to be collected and communicated.
    2. Identify the number of images that need to be taken and communicated.
    3. For each rover, calculate the distance to the nearest waypoint where it can perform the required task.
    4. Sum the estimated actions for each task, including navigation, sampling, imaging, and communication.
    """

    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:
            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])

        # Build a map of waypoints to their visible waypoints
        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])

        # Build a map of rovers to their stores
        self.rover_stores = {}
        for fact in self.static:
            if match(fact, "store_of", "*", "*"):
                parts = get_parts(fact)
                self.rover_stores[parts[2]] = parts[1]

        # Build a map of cameras to their calibration targets
        self.calibration_targets = {}
        for fact in self.static:
            if match(fact, "calibration_target", "*", "*"):
                parts = get_parts(fact)
                self.calibration_targets[parts[1]] = parts[2]

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

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

        # Initialize the heuristic value
        total_cost = 0

        # Check for soil and rock samples that need to be communicated
        for goal in self.goals:
            if match(goal, "communicated_soil_data", "*"):
                waypoint = get_parts(goal)[1]
                if f"(communicated_soil_data {waypoint})" not in state:
                    total_cost += 2  # Sample and communicate
            elif match(goal, "communicated_rock_data", "*"):
                waypoint = get_parts(goal)[1]
                if f"(communicated_rock_data {waypoint})" not in state:
                    total_cost += 2  # Sample and communicate

        # Check for images that need to be communicated
        for goal in self.goals:
            if match(goal, "communicated_image_data", "*", "*"):
                objective, mode = get_parts(goal)[1], get_parts(goal)[2]
                if f"(communicated_image_data {objective} {mode})" not in state:
                    total_cost += 3  # Calibrate, take image, and communicate

        # Estimate navigation costs
        for rover in self.rovers:
            rover_at = None
            for fact in state:
                if match(fact, "at", rover, "*"):
                    rover_at = get_parts(fact)[2]
                    break
            if rover_at:
                # Estimate the cost to reach the nearest waypoint with a task
                total_cost += 1  # Assume at least one navigation action

        return total_cost
