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:
    - Collecting soil and rock samples.
    - Taking images of objectives.
    - Communicating data to the lander.

    # Assumptions
    - Each rover can carry only one sample at a time (soil or rock).
    - Each rover can take images of objectives if equipped with the appropriate camera.
    - Communication with the lander requires the rover to be at a waypoint visible from the lander's location.

    # Heuristic Initialization
    - Extract goal conditions and static facts from the task.
    - Identify the lander's location, waypoints, and objectives.
    - Map rovers to their stores and cameras.

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the current state of each rover:
       - Location.
       - Whether it is carrying a soil or rock sample.
       - Whether it has taken required images.
    2. For each goal condition:
       - If the goal is to communicate soil data:
         - Check if the rover has the soil sample.
         - Estimate the number of actions to navigate to the lander and communicate.
       - If the goal is to communicate rock data:
         - Check if the rover has the rock sample.
         - Estimate the number of actions to navigate to the lander and communicate.
       - If the goal is to communicate image data:
         - Check if the rover has taken the required image.
         - Estimate the number of actions to navigate to the lander and communicate.
    3. Sum the estimated actions for all goal conditions.
    """

    def __init__(self, task):
        """Initialize the heuristic by extracting goal conditions and static facts."""
        self.goals = task.goals
        self.static = task.static

        # Extract lander location
        self.lander_location = None
        for fact in self.static:
            if match(fact, "at_lander", "*", "*"):
                self.lander_location = get_parts(fact)[2]
                break

        # Map rovers to their stores and cameras
        self.rover_stores = {}
        self.rover_cameras = {}
        for fact in self.static:
            if match(fact, "store_of", "*", "*"):
                store, rover = get_parts(fact)[1], get_parts(fact)[2]
                self.rover_stores[rover] = store
            if match(fact, "on_board", "*", "*"):
                camera, rover = get_parts(fact)[1], get_parts(fact)[2]
                if rover not in self.rover_cameras:
                    self.rover_cameras[rover] = []
                self.rover_cameras[rover].append(camera)

    def __call__(self, node):
        """Estimate the number of actions required to achieve the goal state."""
        state = node.state
        total_cost = 0

        # Check each goal condition
        for goal in self.goals:
            predicate, *args = get_parts(goal)
            if predicate == "communicated_soil_data":
                waypoint = args[0]
                # Find a rover that has the soil sample
                rover_with_soil = None
                for fact in state:
                    if match(fact, "have_soil_analysis", "*", waypoint):
                        rover_with_soil = get_parts(fact)[1]
                        break
                if rover_with_soil:
                    # Estimate actions to communicate
                    total_cost += self.estimate_communication_actions(rover_with_soil, state)
                else:
                    # Estimate actions to collect and communicate
                    total_cost += self.estimate_sample_actions("soil", waypoint, state)

            elif predicate == "communicated_rock_data":
                waypoint = args[0]
                # Find a rover that has the rock sample
                rover_with_rock = None
                for fact in state:
                    if match(fact, "have_rock_analysis", "*", waypoint):
                        rover_with_rock = get_parts(fact)[1]
                        break
                if rover_with_rock:
                    # Estimate actions to communicate
                    total_cost += self.estimate_communication_actions(rover_with_rock, state)
                else:
                    # Estimate actions to collect and communicate
                    total_cost += self.estimate_sample_actions("rock", waypoint, state)

            elif predicate == "communicated_image_data":
                objective, mode = args[0], args[1]
                # Find a rover that has the image
                rover_with_image = None
                for fact in state:
                    if match(fact, "have_image", "*", objective, mode):
                        rover_with_image = get_parts(fact)[1]
                        break
                if rover_with_image:
                    # Estimate actions to communicate
                    total_cost += self.estimate_communication_actions(rover_with_image, state)
                else:
                    # Estimate actions to take and communicate the image
                    total_cost += self.estimate_image_actions(objective, mode, state)

        return total_cost

    def estimate_communication_actions(self, rover, state):
        """Estimate the number of actions required for a rover to communicate data."""
        # Check if the rover is already at a waypoint visible from the lander
        rover_location = None
        for fact in state:
            if match(fact, "at", rover, "*"):
                rover_location = get_parts(fact)[2]
                break
        if rover_location and self.is_visible(rover_location, self.lander_location):
            return 1  # Only the communicate action is needed
        else:
            return 2  # Navigate to a visible waypoint and communicate

    def estimate_sample_actions(self, sample_type, waypoint, state):
        """Estimate the number of actions required to collect and communicate a sample."""
        # Find a rover that can collect the sample
        rover = None
        for fact in state:
            if match(fact, "at", "*", waypoint):
                rover = get_parts(fact)[1]
                break
        if rover:
            # Check if the rover is equipped for the sample type
            equipped = False
            if sample_type == "soil":
                for fact in state:
                    if match(fact, "equipped_for_soil_analysis", rover):
                        equipped = True
                        break
            elif sample_type == "rock":
                for fact in state:
                    if match(fact, "equipped_for_rock_analysis", rover):
                        equipped = True
                        break
            if equipped:
                # Estimate actions to collect and communicate
                return 3  # Sample, navigate, communicate
            else:
                # Rover needs to be equipped first
                return 4  # Equip, sample, navigate, communicate
        else:
            # No rover at the waypoint
            return 5  # Navigate, equip, sample, navigate, communicate

    def estimate_image_actions(self, objective, mode, state):
        """Estimate the number of actions required to take and communicate an image."""
        # Find a rover that can take the image
        rover = None
        for fact in state:
            if match(fact, "equipped_for_imaging", "*"):
                rover = get_parts(fact)[1]
                break
        if rover:
            # Check if the rover has a calibrated camera for the objective
            calibrated = False
            for fact in state:
                if match(fact, "calibrated", "*", rover):
                    camera = get_parts(fact)[1]
                    # Check if the camera supports the mode and is calibrated for the objective
                    for static_fact in self.static:
                        if match(static_fact, "calibration_target", camera, objective):
                            calibrated = True
                            break
                    if calibrated:
                        break
            if calibrated:
                # Estimate actions to take and communicate the image
                return 3  # Take image, navigate, communicate
            else:
                # Rover needs to calibrate the camera first
                return 4  # Calibrate, take image, navigate, communicate
        else:
            # No rover equipped for imaging
            return 5  # Equip, calibrate, take image, navigate, communicate

    def is_visible(self, waypoint1, waypoint2):
        """Check if two waypoints are visible from each other."""
        for fact in self.static:
            if match(fact, "visible", waypoint1, waypoint2) or match(fact, "visible", waypoint2, waypoint1):
                return True
        return False
