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 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 correct location before serving.

    # 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 (e.g., which ingredients are available, which sandwiches exist).

    # Step-By-Step Thinking for Computing Heuristic
    1. Identify the number of children who still need to be served.
    2. For each child:
       - If the child is allergic to gluten, ensure that 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 required sandwich, check if the necessary ingredients 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 correct location:
       - Each tray must be moved to the location where the child is waiting.
    6. Sum the total number of actions required to make sandwiches, place them on trays, move trays, and serve the sandwiches.
    """

    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 and their 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 ingredients.
        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 still need to be served.
        unserved_children = {
            child for child in self.allergic_children | self.not_allergic_children
            if f"(served {child})" not in state
        }

        # Count the number of sandwiches that need to be made.
        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 that need to be placed on trays.
        sandwiches_on_trays = sum(1 for fact in state if match(fact, "ontray", "*", "*"))
        sandwiches_to_place = max(0, sandwiches_needed - sandwiches_on_trays)

        # Count the number of trays that need to be moved.
        trays_at_kitchen = sum(1 for fact in state if match(fact, "at", "*", "kitchen"))
        trays_to_move = max(0, len(unserved_children) - trays_at_kitchen

        # Sum the total number of actions.
        total_cost = (
            sandwiches_needed  # Actions to make sandwiches.
            + sandwiches_to_place  # Actions to place sandwiches on trays.
            + trays_to_move  # Actions to move trays.
            + len(unserved_children)  # Actions to serve sandwiches.
        )

        return total_cost
