"""
norvig.planning – AIMA Chapters 10 & 11
=========================================
STRIPS planning and HTN (Hierarchical Task Network) planning.
HTN is a significant gap – no production Python library covers it well.

Classes / functions
-------------------
Action           – STRIPS action (preconditions, add-effects, delete-effects)
STRIPSProblem    – STRIPS planning problem
strips_planner   – forward-state STRIPS BFS planner (AIMA §10.2)
HTNMethod        – HTN method (decomposes a task into subtasks)
HTNProblem       – HTN planning problem
htn_planner      – depth-first HTN planner (AIMA §11.2)

Example – Blocks World
----------------------
>>> move = Action("Move(A,B)", pre={"clear_A","on_table_A","clear_B"},
...               add={"on_A_B"}, delete={"on_table_A","clear_B"})
>>> prob = STRIPSProblem(
...     initial={"clear_A","on_table_A","clear_B","on_table_B"},
...     goal={"on_A_B"},
...     actions=[move])
>>> plan = strips_planner(prob)
"""

from __future__ import annotations
from collections import deque
from typing import Any, Dict, FrozenSet, List, Optional, Set


# ---------------------------------------------------------------------------
# STRIPS (AIMA §10.2)
# ---------------------------------------------------------------------------

class Action:
    """
    STRIPS action with propositional preconditions and effects.

    Parameters
    ----------
    name       : action name string
    pre        : set of propositions that must be true
    add        : propositions added on execution
    delete     : propositions deleted on execution
    """

    def __init__(self, name: str, pre: Set[str], add: Set[str], delete: Set[str]):
        self.name   = name
        self.pre    = frozenset(pre)
        self.add    = frozenset(add)
        self.delete = frozenset(delete)

    def applicable(self, state: FrozenSet[str]) -> bool:
        return self.pre <= state

    def apply(self, state: FrozenSet[str]) -> FrozenSet[str]:
        return (state - self.delete) | self.add

    def __repr__(self):
        return self.name


class STRIPSProblem:
    def __init__(self, initial: Set[str], goal: Set[str], actions: List[Action]):
        self.initial = frozenset(initial)
        self.goal    = frozenset(goal)
        self.actions = actions


def strips_planner(problem: STRIPSProblem, max_steps: int = 1000) -> Optional[List[Action]]:
    """
    Forward-state BFS STRIPS planner (AIMA §10.2).
    Returns list of actions (plan) or None if unsolvable within max_steps.
    """
    initial = problem.initial
    goal    = problem.goal
    queue   = deque([(initial, [])])
    visited : Set[FrozenSet] = {initial}

    while queue:
        state, plan = queue.popleft()
        if goal <= state:
            return plan
        if len(plan) >= max_steps:
            continue
        for action in problem.actions:
            if action.applicable(state):
                next_state = action.apply(state)
                if next_state not in visited:
                    visited.add(next_state)
                    queue.append((next_state, plan + [action]))
    return None


# ---------------------------------------------------------------------------
# HTN – Hierarchical Task Network (AIMA §11.2)
# ---------------------------------------------------------------------------

class HTNMethod:
    """
    An HTN method decomposes a compound task into an ordered list of subtasks.

    Parameters
    ----------
    task       : name of the compound task this method decomposes
    precond    : set of propositions that must hold to use this method
    subtasks   : ordered list of (sub)task names to achieve
    """

    def __init__(self, task: str, precond: Set[str], subtasks: List[str]):
        self.task     = task
        self.precond  = frozenset(precond)
        self.subtasks = list(subtasks)

    def applicable(self, state: FrozenSet[str]) -> bool:
        return self.precond <= state

    def __repr__(self):
        return f"Method({self.task} → {self.subtasks})"


class HTNProblem:
    """
    HTN planning problem.

    Parameters
    ----------
    initial  : initial world state (set of propositions)
    tasks    : top-level task network (ordered list of task names)
    methods  : dict {task_name: [HTNMethod, ...]}
    actions  : dict {action_name: Action}  – primitive tasks
    """

    def __init__(
        self,
        initial: Set[str],
        tasks: List[str],
        methods: Dict[str, List[HTNMethod]],
        actions: Dict[str, Action],
    ):
        self.initial = frozenset(initial)
        self.tasks   = tasks
        self.methods = methods
        self.actions = actions


def htn_planner(
    problem: HTNProblem,
    max_depth: int = 50,
) -> Optional[List[str]]:
    """
    Depth-first HTN planner (AIMA §11.2 / Nau et al. SHOP-style).

    Returns an ordered list of primitive action names or None if no plan found.
    """
    return _htn_search(problem.initial, list(problem.tasks), [], problem, max_depth)


def _htn_search(state, tasks, plan, problem, depth):
    if depth == 0:
        return None
    if not tasks:
        return plan                                   # empty task network → success

    task = tasks[0]
    rest = tasks[1:]

    # Primitive action
    if task in problem.actions:
        action = problem.actions[task]
        if action.applicable(state):
            new_state = action.apply(state)
            result = _htn_search(new_state, rest, plan + [task], problem, depth - 1)
            if result is not None:
                return result
        return None

    # Compound task – try each applicable method
    for method in problem.methods.get(task, []):
        if method.applicable(state):
            new_tasks = method.subtasks + rest
            result = _htn_search(state, new_tasks, plan, problem, depth - 1)
            if result is not None:
                return result

    return None
