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 appropriate tables before serving.

    # Heuristic Initialization
    - Extract the goal conditions (all children must be served).
    - Extract static facts to identify children with gluten allergies and the locations of tables.
    - Identify the available bread and content portions in the kitchen.

    # 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 has a gluten allergy, ensure that a gluten-free sandwich is available.
       - If the child does not have a gluten allergy, 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 and moved to the appropriate tables.
    5. Sum the 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 children with gluten allergies.
        self.allergic_children = {
            get_parts(fact)[1] for fact in self.static if match(fact, "allergic_gluten", "*")
        }

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

        # Extract 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 unserved children.
        unserved_children = {
            child for child in self.waiting_children if f"(served {child})" not in state
        }

        # Initialize the heuristic cost.
        total_cost = 0

        # Count the number of sandwiches that need to be made.
        sandwiches_needed = len(unserved_children)
        gluten_free_sandwiches_needed = sum(
            1 for child in unserved_children if child in self.allergic_children
        )

        # Count the number of available gluten-free and regular sandwiches.
        available_gluten_free_sandwiches = sum(
            1 for fact in state if match(fact, "no_gluten_sandwich", "*")
        )
        available_regular_sandwiches = sum(
            1 for fact in state if match(fact, "at_kitchen_sandwich", "*")
        ) - available_gluten_free_sandwiches

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

        # Add the cost of making sandwiches.
        total_cost += sandwiches_to_make * 1  # Each sandwich requires one action to make.

        # 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)

        # Add the cost of placing sandwiches on trays.
        total_cost += sandwiches_to_place * 1  # Each sandwich requires one action to place on a tray.

        # Count the number of trays that need to be moved.
        trays_to_move = len(set(self.waiting_children.values()))  # Each table requires one tray.

        # Add the cost of moving trays.
        total_cost += trays_to_move * 1  # Each tray requires one action to move.

        # Add the cost of serving sandwiches.
        total_cost += sandwiches_needed * 1  # Each sandwich requires one action to serve.

        return total_cost
