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 tray1 kitchen)".
    - `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 with sandwiches, considering the constraints of gluten allergies and the availability of ingredients.

    # Assumptions
    - Each child must be served exactly one sandwich.
    - Children with gluten allergies must be served gluten-free sandwiches.
    - Sandwiches must be made from available bread and content portions in the kitchen.
    - Sandwiches must be placed on trays and moved to the children's locations.

    # Heuristic Initialization
    - Extract the goal conditions (all children must be served).
    - Extract static facts (e.g., which children are allergic to gluten, which bread and content portions are gluten-free).
    - Identify the initial state of the kitchen (available bread and content portions, trays, and sandwiches).

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the number of children who are still waiting to be served.
    2. For each waiting child:
       - If the child is allergic to gluten, ensure that a gluten-free sandwich is available or can be made.
       - If the child is not allergic to gluten, any sandwich can be used.
    3. Count the number of sandwiches that need to be made:
       - For each required sandwich, check if the necessary bread and content portions 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 to the children's locations:
       - Each tray must be moved to the correct location before the sandwiches can be served.
    6. Sum the total number of actions required to make sandwiches, place them on trays, and move the trays to the children's locations.
    """

    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 information about children's allergies.
        self.allergic_children = {
            get_parts(fact)[1] for fact in self.static if match(fact, "allergic_gluten", "*")
        }
        self.not_allergic_children = {
            get_parts(fact)[1] for fact in self.static if match(fact, "not_allergic_gluten", "*")
        }

        # Extract information about gluten-free bread and content portions.
        self.gluten_free_bread = {
            get_parts(fact)[1] for fact in self.static if match(fact, "no_gluten_bread", "*")
        }
        self.gluten_free_content = {
            get_parts(fact)[1] for fact in self.static if match(fact, "no_gluten_content", "*")
        }

    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 still waiting to be served.
        waiting_children = sum(1 for fact in state if match(fact, "waiting", "*", "*"))

        # Initialize the total cost.
        total_cost = 0

        # For each waiting child, estimate the actions required to serve them.
        for child in self.allergic_children | self.not_allergic_children:
            if f"(waiting {child} *)" in state:
                # Check if the child is allergic to gluten.
                if child in self.allergic_children:
                    # Ensure a gluten-free sandwich is available or can be made.
                    gluten_free_sandwich_available = any(
                        match(fact, "no_gluten_sandwich", "*") for fact in state
                    )
                    if not gluten_free_sandwich_available:
                        # Need to make a gluten-free sandwich.
                        total_cost += 1  # make_sandwich_no_gluten action.
                else:
                    # Any sandwich can be used.
                    sandwich_available = any(
                        match(fact, "at_kitchen_sandwich", "*") for fact in state
                    )
                    if not sandwich_available:
                        # Need to make a sandwich.
                        total_cost += 1  # make_sandwich action.

                # Place the sandwich on a tray.
                total_cost += 1  # put_on_tray action.

                # Move the tray to the child's location.
                total_cost += 1  # move_tray action.

                # Serve the sandwich to the child.
                total_cost += 1  # serve_sandwich or serve_sandwich_no_gluten action.

        return total_cost
