"""
norvig.search – AIMA Chapters 3 & 4
=====================================
Uses networkx as the graph engine.  Two modes:

1. Explicit graph (networkx.DiGraph):
   Pass a pre-built graph and use GraphProblem to wrap it.

2. Abstract Problem:
   Lazily expands into a networkx DiGraph; all search algorithms call
   networkx.astar_path / networkx.dijkstra_path / networkx.bfs_edges
   underneath.

    from ml_ai_library.norvig.search import (
        Problem, Node, GraphProblem,
        astar_search,           # nx.astar_path
        ucs_search,             # nx.dijkstra_path
        bfs_search,             # nx.shortest_path (BFS)
        dfs_search,             # nx.dfs_preorder_nodes
        iddfs_search,           # nx.dfs_preorder_nodes with depth_limit
        greedy_search,          # nx.astar_path (ignore g)
        beam_search,            # bounded BFS via nx successors
    )
"""

from __future__ import annotations
from collections import deque
from typing import Any, Callable, List, Optional

import networkx as nx


# ---------------------------------------------------------------------------
# Abstract problem interface (AIMA §3.3)
# ---------------------------------------------------------------------------

class Problem:
    """AIMA Problem base class.  Subclass and override the five methods."""

    def __init__(self, initial: Any, goal: Any = None):
        self.initial = initial
        self.goal    = goal

    def actions(self, state: Any) -> List:
        raise NotImplementedError

    def result(self, state: Any, action: Any) -> Any:
        raise NotImplementedError

    def goal_test(self, state: Any) -> bool:
        return state == self.goal

    def path_cost(self, cost: float, s1: Any, action: Any, s2: Any) -> float:
        return cost + 1

    def h(self, node: "Node") -> float:
        """Heuristic h(n). Default 0 gives Dijkstra behaviour."""
        return 0.0


class Node:
    """Search tree node; stores state, parent, action and path cost."""

    def __init__(self, state, parent=None, action=None, path_cost: float = 0.0):
        self.state     = state
        self.parent    = parent
        self.action    = action
        self.path_cost = path_cost
        self.depth     = 0 if parent is None else parent.depth + 1

    def solution(self) -> List:
        return [n.action for n in self._path()[1:]]

    def _path(self) -> List["Node"]:
        node, path = self, []
        while node:
            path.append(node)
            node = node.parent
        return list(reversed(path))

    def __repr__(self):
        return f"<Node {self.state}>"

    def __lt__(self, other):
        return self.path_cost < other.path_cost


# ---------------------------------------------------------------------------
# Problem -> networkx DiGraph (lazy BFS expansion used by all algorithms)
# ---------------------------------------------------------------------------

def _expand_to_graph(problem: Problem, max_nodes: int = 50_000) -> nx.DiGraph:
    """Expand `problem` into nx.DiGraph.  Edge attrs: weight, action."""
    G = nx.DiGraph()
    G.add_node(problem.initial)
    frontier: deque = deque([problem.initial])
    visited: set = {problem.initial}
    while frontier and len(G) < max_nodes:
        s = frontier.popleft()
        for a in problem.actions(s):
            s2 = problem.result(s, a)
            w  = problem.path_cost(0, s, a, s2)
            G.add_edge(s, s2, weight=w, action=a)
            if s2 not in visited:
                visited.add(s2)
                frontier.append(s2)
    return G


def _goal_node(problem: Problem, G: nx.DiGraph) -> Optional[Any]:
    return next((n for n in G.nodes if problem.goal_test(n)), None)


def _path_to_node(G: nx.DiGraph, source, target, node_path: List) -> Node:
    """Convert nx node-list to linked Node chain."""
    node = Node(source)
    for i in range(1, len(node_path)):
        edge = G[node_path[i - 1]][node_path[i]]
        node = Node(node_path[i], node, edge.get("action"),
                    node.path_cost + edge.get("weight", 1))
    return node


# ---------------------------------------------------------------------------
# Search algorithms – all delegate to networkx subroutines
# ---------------------------------------------------------------------------

def astar_search(problem: Problem, h: Optional[Callable] = None) -> Optional[Node]:
    """A* via networkx.astar_path (AIMA Fig 3.14)."""
    hfn  = h or problem.h
    G    = _expand_to_graph(problem)
    goal = _goal_node(problem, G)
    if goal is None:
        return None
    try:
        path = nx.astar_path(G, problem.initial, goal,
                             heuristic=lambda u, _v: hfn(Node(u)),
                             weight="weight")
        return _path_to_node(G, problem.initial, goal, path)
    except (nx.NetworkXNoPath, nx.NodeNotFound):
        return None


def ucs_search(problem: Problem) -> Optional[Node]:
    """Uniform-cost search via networkx.dijkstra_path (AIMA Fig 3.14)."""
    G    = _expand_to_graph(problem)
    goal = _goal_node(problem, G)
    if goal is None:
        return None
    try:
        path = nx.dijkstra_path(G, problem.initial, goal, weight="weight")
        return _path_to_node(G, problem.initial, goal, path)
    except (nx.NetworkXNoPath, nx.NodeNotFound):
        return None

uniform_cost_search = ucs_search


def bfs_search(problem: Problem) -> Optional[Node]:
    """BFS via networkx.shortest_path (AIMA Fig 3.11)."""
    G    = _expand_to_graph(problem)
    goal = _goal_node(problem, G)
    if goal is None:
        return None
    try:
        path = nx.shortest_path(G, problem.initial, goal)
        return _path_to_node(G, problem.initial, goal, path)
    except (nx.NetworkXNoPath, nx.NodeNotFound):
        return None

breadth_first_search = bfs_search


def dfs_search(problem: Problem) -> Optional[Node]:
    """DFS via networkx.dfs_preorder_nodes (AIMA §3.4.3)."""
    G = _expand_to_graph(problem)
    for state in nx.dfs_preorder_nodes(G, problem.initial):
        if problem.goal_test(state):
            path = nx.shortest_path(G, problem.initial, state)
            return _path_to_node(G, problem.initial, state, path)
    return None

depth_first_search = dfs_search


def iddfs_search(problem: Problem, max_depth: int = 50) -> Optional[Node]:
    """IDDFS via nx.dfs_preorder_nodes with depth_limit (AIMA Fig 3.18)."""
    G = _expand_to_graph(problem)
    for depth in range(max_depth + 1):
        for state in nx.dfs_preorder_nodes(G, problem.initial, depth_limit=depth):
            if problem.goal_test(state):
                path = nx.shortest_path(G, problem.initial, state)
                return _path_to_node(G, problem.initial, state, path)
    return None

iterative_deepening_search = iddfs_search


def greedy_search(problem: Problem, h: Optional[Callable] = None) -> Optional[Node]:
    """Greedy best-first via nx.astar_path ignoring path cost (AIMA §3.5.1)."""
    hfn  = h or problem.h
    G    = _expand_to_graph(problem)
    goal = _goal_node(problem, G)
    if goal is None:
        return None
    try:
        path = nx.astar_path(G, problem.initial, goal,
                             heuristic=lambda u, _v: hfn(Node(u)),
                             weight=lambda u, v, d: 0)
        return _path_to_node(G, problem.initial, goal, path)
    except (nx.NetworkXNoPath, nx.NodeNotFound):
        return None

greedy_best_first_search = greedy_search


def beam_search(problem: Problem, beam_width: int = 3,
                h: Optional[Callable] = None) -> Optional[Node]:
    """Beam search (AIMA §4.1.2) using networkx for neighbour expansion."""
    hfn = h or problem.h
    G   = _expand_to_graph(problem)
    beam = [problem.initial]
    while beam:
        for s in beam:
            if problem.goal_test(s):
                path = nx.shortest_path(G, problem.initial, s)
                return _path_to_node(G, problem.initial, s, path)
        candidates = list({nb for s in beam for nb in G.successors(s)})
        if not candidates:
            break
        beam = sorted(candidates, key=lambda s: hfn(Node(s)))[:beam_width]
    return None


# ---------------------------------------------------------------------------
# GraphProblem – wrap a networkx.DiGraph as an AIMA Problem
# ---------------------------------------------------------------------------

class GraphProblem(Problem):
    """
    Wrap a pre-built networkx.DiGraph as a Problem.

    Example
    -------
    >>> import networkx as nx
    >>> G = nx.DiGraph()
    >>> G.add_edge('A', 'B', weight=4); G.add_edge('B', 'C', weight=2)
    >>> prob = GraphProblem(G, initial='A', goal='C')
    >>> astar_search(prob).solution()
    ['B', 'C']
    """

    def __init__(self, graph: nx.DiGraph, initial: Any, goal: Any):
        super().__init__(initial, goal)
        self.graph = graph

    def actions(self, state):
        return list(self.graph.successors(state))

    def result(self, state, action):
        return action     # next node IS the action

    def path_cost(self, cost, s1, action, s2):
        return cost + self.graph[s1][s2].get("weight", 1)
