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 (served children) and static facts (allergies, waiting locations, etc.) from the task.
    - Identify the number of children, the number of trays, and the 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 has a gluten allergy, ensure that a gluten-free sandwich is available or can be made.
       - If the child does not have a gluten allergy, any sandwich can be used.
    3. Estimate the number of actions required to make the necessary sandwiches:
       - Each sandwich requires one bread portion and one content portion.
       - Gluten-free sandwiches require specific gluten-free ingredients.
    4. Estimate the number of actions required to place sandwiches on trays:
       - Each sandwich must be placed on a tray, and trays must be moved to the appropriate tables.
    5. Estimate the number of actions required to serve the sandwiches:
       - Each sandwich must be served to the child at their waiting location.
    """

    def __init__(self, task):
        """Initialize the heuristic by extracting goal conditions and static facts."""
        self.goals = task.goals  # Goal conditions.
        static_facts = task.static  # Facts that are not affected by actions.

        # Extract information about children and their allergies.
        self.allergic_children = {
            get_parts(fact)[1] for fact in static_facts if match(fact, "allergic_gluten", "*")
        }
        self.not_allergic_children = {
            get_parts(fact)[1] for fact in static_facts if match(fact, "not_allergic_gluten", "*")
        }

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

        # Extract gluten-free ingredients.
        self.gluten_free_bread = {
            get_parts(fact)[1] for fact in static_facts if match(fact, "no_gluten_bread", "*")
        }
        self.gluten_free_content = {
            get_parts(fact)[1] for fact in static_facts 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 not yet served.
        unserved_children = {
            child for child in self.allergic_children | self.not_allergic_children
            if f"(served {child})" not in state
        }

        total_cost = 0  # Initialize action cost counter.

        for child in unserved_children:
            # Determine if the child has a gluten allergy.
            has_allergy = child in self.allergic_children

            # Estimate the cost of making a sandwich.
            if has_allergy:
                # 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:
                    # Need to make a gluten-free sandwich.
                    total_cost += 1  # Make sandwich 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.

            # Estimate the cost of placing the sandwich on a tray.
            total_cost += 1  # Put on tray action.

            # Estimate the cost of moving the tray to the child's location.
            total_cost += 1  # Move tray action.

            # Estimate the cost of serving the sandwich.
            total_cost += 1  # Serve sandwich action.

        return total_cost
