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 appropriate tables 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 trays are at the kitchen).

    # 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 gluten-free sandwiches, ensure that both gluten-free bread and content portions are available.
       - For regular sandwiches, ensure that any bread and content portions are available.
    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 appropriate tables:
       - Each tray must be moved to the table 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 children.
    """

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

        # Count the number of sandwiches that need to be made.
        sandwiches_to_make = 0
        for child in self.allergic_children:
            if f"(waiting {child} *)" in state:
                # Check if a gluten-free sandwich is already available.
                gluten_free_sandwich_available = any(
                    match(fact, "no_gluten_sandwich", "*") for fact in state
                )
                if not gluten_free_sandwich_available:
                    # Check if gluten-free ingredients are available.
                    gluten_free_bread_available = any(
                        match(fact, "at_kitchen_bread", bread) and bread in self.gluten_free_bread
                        for fact in state
                    )
                    gluten_free_content_available = any(
                        match(fact, "at_kitchen_content", content) and content in self.gluten_free_content
                        for fact in state
                    )
                    if gluten_free_bread_available and gluten_free_content_available:
                        sandwiches_to_make += 1

        for child in self.not_allergic_children:
            if f"(waiting {child} *)" in state:
                # Check if any sandwich is already available.
                sandwich_available = any(match(fact, "at_kitchen_sandwich", "*") for fact in state)
                if not sandwich_available:
                    # Check if any ingredients are available.
                    bread_available = any(match(fact, "at_kitchen_bread", "*") for fact in state)
                    content_available = any(match(fact, "at_kitchen_content", "*") for fact in state)
                    if bread_available and content_available:
                        sandwiches_to_make += 1

        # Count the number of sandwiches that need to be placed on trays.
        sandwiches_to_place = waiting_children

        # Count the number of trays that need to be moved.
        trays_to_move = waiting_children

        # Total heuristic value is the sum of all required actions.
        total_cost = sandwiches_to_make + sandwiches_to_place + trays_to_move

        return total_cost
