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 navigation required for rovers to reach waypoints where samples or images are located.
    - The calibration of cameras before taking images.

    # Assumptions
    - Rovers can only carry one sample at a time (soil or rock).
    - Cameras must be calibrated before taking images.
    - 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 map rovers to their stores, cameras, and waypoints.
    - Identify the lander's location and the visibility graph between waypoints.

    # 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 number of navigation steps required to reach waypoints with samples or objectives.
       - Account for the calibration of cameras if needed.
       - Estimate the number of actions required to collect samples, take images, and communicate data.
    4. Sum the estimated actions across all rovers to compute the total heuristic value.
    """

    def __init__(self, task):
        """Initialize the heuristic by extracting goal conditions and static facts."""
        self.goals = task.goals  # Goal conditions.
        static_facts = task.static  # Facts that are not affected by actions.

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

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

        # Map rovers to their onboard cameras.
        self.rover_cameras = {
            get_parts(fact)[2]: get_parts(fact)[1]
            for fact in static_facts
            if match(fact, "on_board", "*", "*")
        }

        # Identify the lander's location.
        self.lander_location = next(
            get_parts(fact)[2]
            for fact in static_facts
            if match(fact, "at_lander", "*", "*")
        )

        # Build a visibility graph between waypoints.
        self.visibility_graph = {}
        for fact in static_facts:
            if match(fact, "visible", "*", "*"):
                waypoint1, waypoint2 = get_parts(fact)[1], get_parts(fact)[2]
                if waypoint1 not in self.visibility_graph:
                    self.visibility_graph[waypoint1] = set()
                self.visibility_graph[waypoint1].add(waypoint2)

    def __call__(self, node):
        """Estimate the number of actions required to achieve the goal state."""
        state = node.state  # Current world 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 not any(match(fact, "communicated_soil_data", waypoint) for fact in state):
                    total_cost += 2  # Collect and communicate soil sample.
            elif match(goal, "communicated_rock_data", "*"):
                waypoint = get_parts(goal)[1]
                if not any(match(fact, "communicated_rock_data", waypoint) for fact in state):
                    total_cost += 2  # Collect and communicate rock sample.

        # 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 not any(match(fact, "communicated_image_data", objective, mode) for fact in state):
                    total_cost += 3  # Calibrate, take image, and communicate.

        # Add navigation costs for rovers to reach waypoints.
        for rover in self.rover_stores:
            rover_location = next(
                get_parts(fact)[2]
                for fact in state
                if match(fact, "at", rover, "*")
            )
            if rover_location != self.lander_location:
                total_cost += 1  # Navigate to lander's location.

        return total_cost
