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_kitchen_bread bread1)".
    - `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 ChildsnackHeuristic(Heuristic):
    """
    A domain-dependent heuristic for the Childsnacks domain.

    # Summary
    This heuristic estimates the number of actions required to serve all children by:
    - Counting the number of sandwiches that need to be made.
    - Counting the number of sandwiches that need to be placed on trays.
    - Counting the number of trays that need to be moved to serve the children.
    - Considering the constraints for gluten-free sandwiches for allergic children.

    # Assumptions
    - Each child requires exactly one sandwich.
    - Gluten-free sandwiches must be served to allergic children.
    - Sandwiches can only be made in the kitchen.
    - Trays must be moved to the correct location to serve sandwiches.

    # Heuristic Initialization
    - Extract the goal conditions (served children) and static facts (allergic children, waiting locations, etc.).
    - Identify the number of children and their requirements (gluten-free or not).

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the number of children who are not yet served.
    2. For each unserved child:
       - If the child is allergic to gluten, ensure a gluten-free sandwich is available.
       - If the child is not allergic, any sandwich can be served.
    3. Count the number of sandwiches that need to be made:
       - For each required sandwich, check if the necessary bread and content are available in the kitchen.
    4. Count the number of sandwiches that need to be placed on trays:
       - Each sandwich must be placed on a tray before it can be served.
    5. Count the number of trays that need to be moved:
       - Trays must be moved to the correct location to serve the sandwiches.
    6. Sum the actions required for making sandwiches, placing them on trays, and moving trays.
    """

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

        # Extract allergic and non-allergic children from static facts.
        self.allergic_children = {
            get_parts(fact)[1] for fact in self.static if match(fact, "allergic_gluten", "*")
        }
        self.non_allergic_children = {
            get_parts(fact)[1] for fact in self.static if match(fact, "not_allergic_gluten", "*")
        }

        # Extract waiting locations for children.
        self.waiting_locations = {
            get_parts(fact)[1]: get_parts(fact)[2]
            for fact in self.static
            if match(fact, "waiting", "*", "*")
        }

    def __call__(self, node):
        """Estimate the number of actions required to serve all children."""
        state = node.state  # Current world state.

        # Count the number of children who are not yet served.
        unserved_children = {
            child for child in self.allergic_children | self.non_allergic_children
            if f"(served {child})" not in state
        }

        # Initialize action cost counters.
        make_sandwich_cost = 0
        put_on_tray_cost = 0
        move_tray_cost = 0

        # For each unserved child, determine the required actions.
        for child in unserved_children:
            # Check if the child is allergic to gluten.
            is_allergic = child in self.allergic_children

            # Determine if a suitable sandwich is already on a tray.
            suitable_sandwich_on_tray = False
            for fact in state:
                if match(fact, "ontray", "*", "*"):
                    sandwich, tray = get_parts(fact)[1], get_parts(fact)[2]
                    if is_allergic and f"(no_gluten_sandwich {sandwich})" in state:
                        suitable_sandwich_on_tray = True
                        break
                    elif not is_allergic:
                        suitable_sandwich_on_tray = True
                        break

            if not suitable_sandwich_on_tray:
                # A sandwich needs to be made and placed on a tray.
                make_sandwich_cost += 1
                put_on_tray_cost += 1

            # Check if the tray is at the correct location.
            tray_at_correct_location = False
            for fact in state:
                if match(fact, "at", "*", self.waiting_locations[child]):
                    tray_at_correct_location = True
                    break

            if not tray_at_correct_location:
                # The tray needs to be moved.
                move_tray_cost += 1

        # Sum the total cost.
        total_cost = make_sandwich_cost + put_on_tray_cost + move_tray_cost

        return total_cost
