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 (served children) and static facts (allergies, waiting locations, etc.).
    - Identify the number of children, trays, and available ingredients.

    # Step-By-Step Thinking for Computing Heuristic
    1. Count 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 used.
    3. Count the number of sandwiches that need to be made:
       - For each unserved child, check if a suitable sandwich is already on a tray or needs to be made.
    4. Count the number of trays that need to be moved to serve the sandwiches:
       - Each tray can carry multiple sandwiches, but it must be moved to the child's location.
    5. Sum the actions required:
       - Making sandwiches (1 action per sandwich).
       - Placing sandwiches on trays (1 action per sandwich).
       - Moving trays to the children's locations (1 action per tray).
       - Serving sandwiches (1 action per child).
    """

    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 the list of children and their allergies.
        self.children = set()
        self.allergic_children = set()
        self.not_allergic_children = set()
        for fact in self.static:
            if match(fact, "allergic_gluten", "*"):
                self.allergic_children.add(get_parts(fact)[1])
                self.children.add(get_parts(fact)[1])
            elif match(fact, "not_allergic_gluten", "*"):
                self.not_allergic_children.add(get_parts(fact)[1])
                self.children.add(get_parts(fact)[1])

        # Extract the list of trays.
        self.trays = set()
        for fact in self.static:
            if match(fact, "at", "tray*", "*"):
                self.trays.add(get_parts(fact)[1])

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

        # Count the number of unserved children.
        unserved_children = set()
        for child in self.children:
            if f"(served {child})" not in state:
                unserved_children.add(child)

        # Count the number of sandwiches needed.
        sandwiches_needed = 0
        for child in unserved_children:
            if child in self.allergic_children:
                # Need a gluten-free sandwich.
                sandwiches_needed += 1
            else:
                # Any sandwich will do.
                sandwiches_needed += 1

        # Count the number of sandwiches already on trays.
        sandwiches_on_trays = 0
        for fact in state:
            if match(fact, "ontray", "*", "*"):
                sandwiches_on_trays += 1

        # Calculate the number of sandwiches that need to be made.
        sandwiches_to_make = max(0, sandwiches_needed - sandwiches_on_trays)

        # Calculate the number of trays that need to be moved.
        trays_to_move = 0
        for tray in self.trays:
            if f"(at {tray} kitchen)" in state:
                trays_to_move += 1

        # Sum the actions required.
        total_cost = (
            sandwiches_to_make  # Make sandwiches.
            + sandwiches_needed  # Place sandwiches on trays.
            + trays_to_move  # Move trays to children's locations.
            + len(unserved_children)  # Serve sandwiches.
        )

        return total_cost
