// Copyright 2026 Gregorio Momm
// Licensed under the Apache License, Version 2.0
//
// Ported from RustworkxCore (IBM Qiskit) — Apache License 2.0
// Source: https://github.com/Qiskit/rustworkx/tree/main/rustworkx-core/src/traversal

//! Graph traversal algorithms ported from RustworkxCore.
//!
//! - `bfs_layers` — iterate nodes layer by layer from a source
//! - `dfs_edges` — DFS edge enumeration (tree/back/forward/cross)
//! - `ancestors` — all nodes reachable in reverse from a node
//! - `descendants` — all nodes reachable forward from a node
//! - `bfs_successors` — (parent, Vec<children>) pairs in BFS order
//! - `bfs_predecessors` — (child, Vec<parents>) pairs in BFS order

use std::collections::{HashSet, VecDeque};

use super::super::graph::{Graph, NodeId};

// ─── BFS Layers ──────────────────────────────────────────────────────────────

/// Return nodes grouped by BFS layer from `source`.
///
/// Layer 0 = [source], layer 1 = direct neighbors, etc.
///
/// Runtime: O(n + m)
pub fn bfs_layers(graph: &Graph, source: NodeId) -> Vec<Vec<NodeId>> {
    let n = graph.upper_node_id_bound() as usize;
    let mut visited = vec![false; n];
    let mut layers: Vec<Vec<NodeId>> = Vec::new();

    visited[source as usize] = true;
    let mut current_layer = vec![source];

    while !current_layer.is_empty() {
        let mut next_layer = Vec::new();
        for &node in &current_layer {
            for nbr in graph.out_neighbors(node).iter().map(|e| e.target) {
                if !visited[nbr as usize] {
                    visited[nbr as usize] = true;
                    next_layer.push(nbr);
                }
            }
        }
        layers.push(current_layer);
        current_layer = next_layer;
    }

    layers
}

// ─── DFS Edges ────────────────────────────────────────────────────────────────

/// Type of a DFS edge.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DfsEdgeType {
    /// Tree edge: leads to an unvisited node
    Tree,
    /// Back edge: leads to an ancestor (cycle)
    Back,
    /// Forward edge: leads to a descendant not via tree
    Forward,
    /// Cross edge: leads to a node in a different branch
    Cross,
}

/// A directed edge seen during DFS.
#[derive(Debug, Clone)]
pub struct DfsEdge {
    pub u: NodeId,
    pub v: NodeId,
    pub edge_type: DfsEdgeType,
}

/// Enumerate edges encountered during DFS from every unvisited node.
///
/// Returns all edges with their type (tree / back / forward / cross).
///
/// Runtime: O(n + m)
pub fn dfs_edges(graph: &Graph) -> Vec<DfsEdge> {
    let n = graph.upper_node_id_bound() as usize;
    // 0 = unvisited, 1 = in stack (grey), 2 = done (black)
    let mut color = vec![0u8; n];
    // Discovery time
    let mut disc = vec![0u32; n];
    let mut timer = 0u32;
    let mut result = Vec::new();

    for start in graph.nodes() {
        if color[start as usize] != 0 {
            continue;
        }

        // Iterative DFS
        let mut stack: Vec<(NodeId, usize)> = vec![(start, 0)];
        color[start as usize] = 1;
        disc[start as usize] = timer;
        timer += 1;

        while let Some((node, idx)) = stack.last_mut() {
            let node = *node;
            let neighbors: Vec<NodeId> = graph.out_neighbors(node).iter().map(|e| e.target).collect();

            if *idx < neighbors.len() {
                let next = neighbors[*idx];
                *idx += 1;

                let edge_type = match color[next as usize] {
                    0 => {
                        // Unvisited → tree edge
                        color[next as usize] = 1;
                        disc[next as usize] = timer;
                        timer += 1;
                        stack.push((next, 0));
                        DfsEdgeType::Tree
                    }
                    1 => DfsEdgeType::Back, // In stack → back edge (cycle)
                    _ => {
                        // Already done
                        if disc[next as usize] > disc[node as usize] {
                            DfsEdgeType::Forward
                        } else {
                            DfsEdgeType::Cross
                        }
                    }
                };
                result.push(DfsEdge { u: node, v: next, edge_type });
            } else {
                color[node as usize] = 2;
                stack.pop();
            }
        }
    }

    result
}

// ─── Ancestors & Descendants ─────────────────────────────────────────────────

/// Return all ancestors of `node` (nodes that can reach `node`).
///
/// Performs BFS/DFS on reversed edges.
///
/// Runtime: O(n + m)
pub fn ancestors(graph: &Graph, node: NodeId) -> HashSet<NodeId> {
    let n = graph.upper_node_id_bound() as usize;
    let mut visited = vec![false; n];
    let mut result = HashSet::new();

    // Build reverse adjacency
    let mut reverse: Vec<Vec<NodeId>> = vec![Vec::new(); n];
    for u in graph.nodes() {
        for e in graph.out_neighbors(u) {
            reverse[e.target as usize].push(u);
        }
    }

    let mut queue = VecDeque::new();
    visited[node as usize] = true;
    queue.push_back(node);

    while let Some(cur) = queue.pop_front() {
        for &pred in &reverse[cur as usize] {
            if !visited[pred as usize] {
                visited[pred as usize] = true;
                result.insert(pred);
                queue.push_back(pred);
            }
        }
    }

    result
}

/// Return all descendants of `node` (nodes reachable from `node`).
///
/// Performs BFS on forward edges.
///
/// Runtime: O(n + m)
pub fn descendants(graph: &Graph, node: NodeId) -> HashSet<NodeId> {
    let n = graph.upper_node_id_bound() as usize;
    let mut visited = vec![false; n];
    let mut result = HashSet::new();

    let mut queue = VecDeque::new();
    visited[node as usize] = true;
    queue.push_back(node);

    while let Some(cur) = queue.pop_front() {
        for nbr in graph.out_neighbors(cur).iter().map(|e| e.target) {
            if !visited[nbr as usize] {
                visited[nbr as usize] = true;
                result.insert(nbr);
                queue.push_back(nbr);
            }
        }
    }

    result
}

// ─── BFS Successors / Predecessors ───────────────────────────────────────────

/// BFS successors: returns (node, children) pairs in BFS order from `source`.
///
/// Each pair gives a node and the list of its children in the BFS tree.
///
/// Runtime: O(n + m)
pub fn bfs_successors(graph: &Graph, source: NodeId) -> Vec<(NodeId, Vec<NodeId>)> {
    let n = graph.upper_node_id_bound() as usize;
    let mut visited = vec![false; n];
    let mut result = Vec::new();

    visited[source as usize] = true;
    let mut queue = VecDeque::new();
    queue.push_back(source);

    while let Some(node) = queue.pop_front() {
        let mut children = Vec::new();
        for nbr in graph.out_neighbors(node).iter().map(|e| e.target) {
            if !visited[nbr as usize] {
                visited[nbr as usize] = true;
                children.push(nbr);
                queue.push_back(nbr);
            }
        }
        if !children.is_empty() {
            result.push((node, children));
        }
    }

    result
}

/// BFS predecessors: returns (node, parents) pairs in BFS order from `source`.
///
/// Each pair gives a node and the list of nodes that first discovered it.
///
/// Runtime: O(n + m)
pub fn bfs_predecessors(graph: &Graph, source: NodeId) -> Vec<(NodeId, Vec<NodeId>)> {
    let n = graph.upper_node_id_bound() as usize;
    let mut visited = vec![false; n];
    let mut result = Vec::new();

    visited[source as usize] = true;
    let mut queue = VecDeque::new();
    queue.push_back(source);

    while let Some(node) = queue.pop_front() {
        for nbr in graph.out_neighbors(node).iter().map(|e| e.target) {
            if !visited[nbr as usize] {
                visited[nbr as usize] = true;
                // Node `nbr` was discovered from `node`
                result.push((nbr, vec![node]));
                queue.push_back(nbr);
            }
        }
    }

    result
}

// ─── Tests ────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::graph::backends::networkit_rust::graph::GraphConfig;

    fn path_directed(n: usize) -> Graph {
        let mut g = Graph::new(GraphConfig::directed());
        for _ in 0..n { g.add_node(); }
        for i in 0..(n as NodeId - 1) { g.add_edge(i, i + 1, None); }
        g
    }

    fn path_undirected(n: usize) -> Graph {
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..n { g.add_node(); }
        for i in 0..(n as NodeId - 1) { g.add_edge(i, i + 1, None); }
        g
    }

    // ── bfs_layers ────────────────────────────────────────────────────────────

    #[test]
    fn test_bfs_layers_path() {
        let g = path_directed(4); // 0→1→2→3
        let layers = bfs_layers(&g, 0);
        assert_eq!(layers.len(), 4);
        assert_eq!(layers[0], vec![0]);
        assert_eq!(layers[1], vec![1]);
        assert_eq!(layers[2], vec![2]);
        assert_eq!(layers[3], vec![3]);
    }

    #[test]
    fn test_bfs_layers_star() {
        let mut g = Graph::new(GraphConfig::directed());
        for _ in 0..5 { g.add_node(); }
        for i in 1..5u64 { g.add_edge(0, i, None); }
        let layers = bfs_layers(&g, 0);
        assert_eq!(layers.len(), 2);
        assert_eq!(layers[0], vec![0]);
        assert_eq!(layers[1].len(), 4);
    }

    // ── dfs_edges ─────────────────────────────────────────────────────────────

    #[test]
    fn test_dfs_edges_path() {
        let g = path_directed(4); // 0→1→2→3
        let edges = dfs_edges(&g);
        // All edges in a path should be tree edges
        for e in &edges {
            assert_eq!(e.edge_type, DfsEdgeType::Tree, "expected tree edge, got {:?}", e.edge_type);
        }
        assert_eq!(edges.len(), 3);
    }

    #[test]
    fn test_dfs_edges_cycle() {
        // 0→1→2→0 creates a back edge
        let mut g = Graph::new(GraphConfig::directed());
        for _ in 0..3 { g.add_node(); }
        g.add_edge(0, 1, None);
        g.add_edge(1, 2, None);
        g.add_edge(2, 0, None);
        let edges = dfs_edges(&g);
        let has_back = edges.iter().any(|e| e.edge_type == DfsEdgeType::Back);
        assert!(has_back, "cycle should produce at least one back edge");
    }

    // ── ancestors & descendants ───────────────────────────────────────────────

    #[test]
    fn test_descendants_path() {
        let g = path_directed(4); // 0→1→2→3
        let desc = descendants(&g, 1);
        assert!(desc.contains(&2));
        assert!(desc.contains(&3));
        assert!(!desc.contains(&0));
        assert!(!desc.contains(&1));
    }

    #[test]
    fn test_ancestors_path() {
        let g = path_directed(4); // 0→1→2→3
        let anc = ancestors(&g, 2);
        assert!(anc.contains(&0));
        assert!(anc.contains(&1));
        assert!(!anc.contains(&2));
        assert!(!anc.contains(&3));
    }

    #[test]
    fn test_ancestors_no_predecessors() {
        let g = path_directed(4); // 0→1→2→3
        let anc = ancestors(&g, 0);
        assert!(anc.is_empty());
    }

    // ── bfs_successors & bfs_predecessors ────────────────────────────────────

    #[test]
    fn test_bfs_successors_path() {
        let g = path_undirected(4); // 0-1-2-3
        let succ = bfs_successors(&g, 0);
        // (0, [1]), (1, [2]), (2, [3])
        assert_eq!(succ.len(), 3);
        assert_eq!(succ[0], (0, vec![1]));
        assert_eq!(succ[1], (1, vec![2]));
        assert_eq!(succ[2], (2, vec![3]));
    }

    #[test]
    fn test_bfs_predecessors_path() {
        let g = path_undirected(4); // 0-1-2-3
        let preds = bfs_predecessors(&g, 0);
        // Each node is discovered from the previous
        assert_eq!(preds.len(), 3);
        assert_eq!(preds[0], (1, vec![0]));
        assert_eq!(preds[1], (2, vec![1]));
        assert_eq!(preds[2], (3, vec![2]));
    }

    #[test]
    fn test_bfs_successors_star() {
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..5 { g.add_node(); }
        for i in 1..5u64 { g.add_edge(0, i, None); }
        let succ = bfs_successors(&g, 0);
        // Root has 4 children, no further children
        assert_eq!(succ.len(), 1);
        assert_eq!(succ[0].0, 0);
        assert_eq!(succ[0].1.len(), 4);
    }
}

// ─── Visitor Pattern ──────────────────────────────────────────────────────────

/// Control flow returned by visitor callbacks.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Control {
    /// Continue traversal normally.
    Continue,
    /// Stop the traversal immediately.
    Break,
    /// Prune: do not traverse into the children of this node.
    Prune,
}

/// BFS traversal with visitor callbacks.
///
/// Calls `discover_vertex(node)` the first time each node is seen,
/// and `examine_edge(u, v)` for every edge examined (whether or not v is new).
///
/// Traversal stops early if any callback returns `Control::Break`.
/// Returning `Control::Prune` from `discover_vertex` prevents expanding that node.
///
/// All nodes reachable from `source` are visited.
pub fn bfs_search<FV, FE>(
    graph: &Graph,
    source: NodeId,
    mut discover_vertex: FV,
    mut examine_edge: FE,
) where
    FV: FnMut(NodeId) -> Control,
    FE: FnMut(NodeId, NodeId) -> Control,
{
    let n = graph.upper_node_id_bound() as usize;
    let mut visited = vec![false; n];
    let mut queue = VecDeque::new();

    visited[source as usize] = true;
    match discover_vertex(source) {
        Control::Break | Control::Prune => return,
        Control::Continue => {}
    }
    queue.push_back(source);

    'outer: while let Some(u) = queue.pop_front() {
        for e in graph.out_neighbors(u).iter() {
            let v = e.target;
            match examine_edge(u, v) {
                Control::Break => break 'outer,
                Control::Prune => continue,
                Control::Continue => {}
            }
            if !visited[v as usize] {
                visited[v as usize] = true;
                match discover_vertex(v) {
                    Control::Break => break 'outer,
                    Control::Prune => {} // don't enqueue
                    Control::Continue => queue.push_back(v),
                }
            }
        }
    }
}

/// DFS traversal with visitor callbacks.
///
/// Calls:
/// - `discover_vertex(node)` when a node is first entered
/// - `finish_vertex(node)` when all children of a node have been explored
/// - `examine_edge(u, v)` for every edge examined
///
/// Traversal stops early if any callback returns `Control::Break`.
/// `Control::Prune` from `discover_vertex` prevents recursing into children.
pub fn dfs_search<FD, FF, FE>(
    graph: &Graph,
    source: NodeId,
    mut discover_vertex: FD,
    mut finish_vertex: FF,
    mut examine_edge: FE,
) where
    FD: FnMut(NodeId) -> Control,
    FF: FnMut(NodeId) -> Control,
    FE: FnMut(NodeId, NodeId) -> Control,
{
    let n = graph.upper_node_id_bound() as usize;
    let mut visited = vec![false; n];
    // Stack items: (node, child_index, pruned)
    let mut stack: Vec<(NodeId, usize, bool)> = Vec::new();

    visited[source as usize] = true;
    let prune = match discover_vertex(source) {
        Control::Break => return,
        Control::Prune => true,
        Control::Continue => false,
    };
    stack.push((source, 0, prune));

    'outer: while let Some((u, child_idx, pruned)) = stack.last_mut() {
        let u = *u;
        let pruned = *pruned;
        let neighbors: Vec<NodeId> = graph.out_neighbors(u).iter().map(|e| e.target).collect();

        if pruned || *child_idx >= neighbors.len() {
            stack.pop();
            match finish_vertex(u) {
                Control::Break => break 'outer,
                _ => {}
            }
            continue;
        }

        let v = neighbors[*child_idx];
        *child_idx += 1;

        match examine_edge(u, v) {
            Control::Break => { break 'outer; }
            Control::Prune => { continue; }
            Control::Continue => {}
        }

        if !visited[v as usize] {
            visited[v as usize] = true;
            let prune_v = match discover_vertex(v) {
                Control::Break => { break 'outer; }
                Control::Prune => true,
                Control::Continue => false,
            };
            stack.push((v, 0, prune_v));
        }
    }
}

/// Dijkstra search with visitor callbacks.
///
/// Processes nodes in order of increasing distance from `source`.
/// Calls `examine_vertex(node, distance)` when a node is settled,
/// and `examine_edge(u, v, weight)` for every relaxed edge.
///
/// Stops early on `Control::Break`. Only works correctly for non-negative weights.
pub fn dijkstra_search<FV, FE>(
    graph: &Graph,
    source: NodeId,
    mut examine_vertex: FV,
    mut examine_edge: FE,
) where
    FV: FnMut(NodeId, f64) -> Control,
    FE: FnMut(NodeId, NodeId, f64) -> Control,
{
    use std::collections::BinaryHeap;
    use std::cmp::Reverse;

    let n = graph.upper_node_id_bound() as usize;
    let mut dist = vec![f64::INFINITY; n];
    let mut settled = vec![false; n];
    dist[source as usize] = 0.0;

    // (Reverse(dist), node)
    let mut heap = BinaryHeap::new();
    heap.push((Reverse(ordered_float(0.0)), source));

    while let Some((Reverse(d_ord), u)) = heap.pop() {
        if settled[u as usize] { continue; }
        let d = f64::from_bits(d_ord);
        settled[u as usize] = true;

        match examine_vertex(u, d) {
            Control::Break => return,
            Control::Prune => continue,
            Control::Continue => {}
        }

        for e in graph.out_neighbors(u).iter() {
            let v = e.target;
            let w = e.weight.unwrap_or(1.0) as f64;
            let nd = d + w;

            match examine_edge(u, v, w) {
                Control::Break => return,
                Control::Prune => continue,
                Control::Continue => {}
            }

            if nd < dist[v as usize] {
                dist[v as usize] = nd;
                heap.push((Reverse(ordered_float(nd)), v));
            }
        }
    }
}

/// Convert f64 to u64 bits preserving total order for non-negative values.
#[inline]
fn ordered_float(v: f64) -> u64 { v.to_bits() }

#[cfg(test)]
mod visitor_tests {
    use super::*;
    use crate::graph::backends::networkit_rust::graph::GraphConfig;

    fn path4() -> Graph {
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..4 { g.add_node(); }
        for i in 0..3u64 { g.add_edge(i, i+1, None); }
        g
    }

    fn star() -> Graph {
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..5 { g.add_node(); }
        for i in 1..5u64 { g.add_edge(0, i, None); }
        g
    }

    // ── bfs_search ────────────────────────────────────────────────────────────

    #[test]
    fn test_bfs_search_visits_all() {
        let g = path4();
        let mut visited = Vec::new();
        bfs_search(&g, 0,
            |n| { visited.push(n); Control::Continue },
            |_, _| Control::Continue,
        );
        let mut v = visited.clone(); v.sort();
        assert_eq!(v, vec![0, 1, 2, 3]);
    }

    #[test]
    fn test_bfs_search_break() {
        let g = path4();
        let mut count = 0;
        bfs_search(&g, 0,
            |_| { count += 1; if count == 2 { Control::Break } else { Control::Continue } },
            |_, _| Control::Continue,
        );
        assert_eq!(count, 2, "should stop after 2 vertices");
    }

    #[test]
    fn test_bfs_search_edge_callback() {
        let g = path4();
        let mut edges = Vec::new();
        bfs_search(&g, 0,
            |_| Control::Continue,
            |u, v| { if u < v { edges.push((u, v)); } Control::Continue },
        );
        assert!(!edges.is_empty());
    }

    // ── dfs_search ────────────────────────────────────────────────────────────

    #[test]
    fn test_dfs_search_visits_all() {
        let g = path4();
        let mut visited = Vec::new();
        dfs_search(&g, 0,
            |n| { visited.push(n); Control::Continue },
            |_| Control::Continue,
            |_, _| Control::Continue,
        );
        let mut v = visited.clone(); v.sort();
        assert_eq!(v, vec![0, 1, 2, 3]);
    }

    #[test]
    fn test_dfs_search_finish_order() {
        // In a path, finish order should be reverse of discovery for the chain
        let g = path4();
        let mut finish_order = Vec::new();
        dfs_search(&g, 0,
            |_| Control::Continue,
            |n| { finish_order.push(n); Control::Continue },
            |_, _| Control::Continue,
        );
        // Last discovered (leaf) finishes first
        assert_eq!(finish_order[0], 3);
    }

    #[test]
    fn test_dfs_search_prune() {
        // Star: prune exploration from center → only center visited
        let g = star();
        let mut visited = Vec::new();
        dfs_search(&g, 0,
            |n| { visited.push(n); Control::Prune }, // prune immediately
            |_| Control::Continue,
            |_, _| Control::Continue,
        );
        assert_eq!(visited, vec![0]);
    }

    // ── dijkstra_search ───────────────────────────────────────────────────────

    #[test]
    fn test_dijkstra_search_order() {
        // In a path, nodes should be settled in order 0,1,2,3
        let g = path4();
        let mut order = Vec::new();
        dijkstra_search(&g, 0,
            |n, _| { order.push(n); Control::Continue },
            |_, _, _| Control::Continue,
        );
        assert_eq!(order, vec![0, 1, 2, 3]);
    }

    #[test]
    fn test_dijkstra_search_break() {
        // Stop as soon as node 2 is settled
        let g = path4();
        let mut settled = Vec::new();
        dijkstra_search(&g, 0,
            |n, _| {
                settled.push(n);
                if n == 2 { Control::Break } else { Control::Continue }
            },
            |_, _, _| Control::Continue,
        );
        assert!(settled.contains(&2));
        assert!(!settled.contains(&3), "should not settle 3 after breaking at 2");
    }
}
