// Copyright 2026 Gregorio Momm
// Licensed under the Apache License, Version 2.0
//
// Ported from RustworkxCore (IBM Qiskit) — Apache License 2.0
//
// Academic References:
// - Cooper, K., Harvey, T., & Kennedy, K. (2001). "A simple, fast dominance
//   algorithm." Software Practice & Experience.

//! Dominance algorithms for directed graphs.
//!
//! - `immediate_dominators` — compute the immediate dominator tree (Cooper et al.)
//! - `dominance_frontiers`  — compute dominance frontiers from the dominator tree

use super::super::graph::{Graph, NodeId};
use std::collections::{HashMap, HashSet};

// ─── Immediate Dominators ─────────────────────────────────────────────────────

/// Compute the immediate dominator of every reachable node in a directed graph.
///
/// Uses the simple iterative dataflow algorithm (Cooper, Harvey & Kennedy 2001),
/// which is efficient in practice even though theoretically O(n²).
///
/// `start` is the entry node (root of the dominator tree).
///
/// Returns a map `node → immediate_dominator`.  The entry node maps to itself.
/// Nodes unreachable from `start` are absent from the map.
///
/// # Panics
/// Does not panic; returns an empty map if `start` is not in the graph.
pub fn immediate_dominators(graph: &Graph, start: NodeId) -> HashMap<NodeId, NodeId> {
    // BFS to get RPO (reverse post-order) and check reachability
    let bound = graph.upper_node_id_bound() as usize;
    let mut visited = vec![false; bound];
    let mut post_order: Vec<NodeId> = Vec::new();
    dfs_post_order(graph, start, &mut visited, &mut post_order);

    if post_order.is_empty() {
        return HashMap::new();
    }

    // Reverse post-order (RPO): entry first
    let mut rpo: Vec<NodeId> = post_order.clone();
    rpo.reverse();

    // RPO index for each node
    let mut rpo_idx = vec![usize::MAX; bound];
    for (i, &n) in rpo.iter().enumerate() {
        rpo_idx[n as usize] = i;
    }

    // Build predecessor list (only for reachable nodes)
    let mut preds: Vec<Vec<NodeId>> = vec![Vec::new(); bound];
    for &u in &rpo {
        for e in graph.out_neighbors(u).iter() {
            if rpo_idx[e.target as usize] != usize::MAX {
                preds[e.target as usize].push(u);
            }
        }
    }

    // idom array: idom[n] = immediate dominator of n (stored as RPO index)
    let mut idom = vec![usize::MAX; bound];
    idom[start as usize] = rpo_idx[start as usize]; // start dominates itself

    let mut changed = true;
    while changed {
        changed = false;
        // Iterate in RPO order (skip entry node at rpo[0])
        for &b in rpo.iter().skip(1) {
            // Find a processed predecessor
            let new_idom_opt = preds[b as usize]
                .iter()
                .find(|&&p| idom[p as usize] != usize::MAX)
                .copied();

            let mut new_idom = match new_idom_opt {
                Some(p) => rpo_idx[p as usize],
                None => continue,
            };

            for &p in &preds[b as usize] {
                if idom[p as usize] != usize::MAX {
                    new_idom = intersect(&idom, new_idom, rpo_idx[p as usize]);
                }
            }

            if idom[b as usize] != new_idom {
                idom[b as usize] = new_idom;
                changed = true;
            }
        }
    }

    // Build output map: NodeId → NodeId
    let mut result = HashMap::new();
    for &n in &rpo {
        let dom_rpo = idom[n as usize];
        if dom_rpo != usize::MAX {
            result.insert(n, rpo[dom_rpo]);
        }
    }
    result
}

/// Intersect two dominators in the RPO lattice (LCA in dominator tree).
fn intersect(idom: &[usize], mut a: usize, mut b: usize) -> usize {
    while a != b {
        while a > b { a = idom[a]; }  // walk up from deeper finger
        while b > a { b = idom[b]; }
    }
    a
}

/// Iterative DFS post-order traversal.
fn dfs_post_order(graph: &Graph, start: NodeId, visited: &mut Vec<bool>, order: &mut Vec<NodeId>) {
    let mut stack: Vec<(NodeId, usize)> = vec![(start, 0)];
    visited[start as usize] = true;

    while let Some((node, child_idx)) = stack.last_mut() {
        let node = *node;
        let neighbors: Vec<NodeId> = graph.out_neighbors(node).iter().map(|e| e.target).collect();
        if *child_idx < neighbors.len() {
            let next = neighbors[*child_idx];
            *child_idx += 1;
            if !visited[next as usize] {
                visited[next as usize] = true;
                stack.push((next, 0));
            }
        } else {
            order.push(node);
            stack.pop();
        }
    }
}

// ─── Dominance Frontiers ──────────────────────────────────────────────────────

/// Compute the dominance frontier of every node.
///
/// The dominance frontier of node `x` is the set of nodes `y` such that
/// `x` dominates a predecessor of `y` but does not strictly dominate `y`.
///
/// Requires a pre-computed `idom` map from `immediate_dominators`.
///
/// Runtime: O(n + m)
pub fn dominance_frontiers(
    graph: &Graph,
    idom: &HashMap<NodeId, NodeId>,
) -> HashMap<NodeId, HashSet<NodeId>> {
    let mut frontiers: HashMap<NodeId, HashSet<NodeId>> =
        idom.keys().map(|&n| (n, HashSet::new())).collect();

    for &b in idom.keys() {
        let preds: Vec<NodeId> = graph
            .nodes()
            .filter(|&u| graph.out_neighbors(u).iter().any(|e| e.target == b))
            .collect();

        if preds.len() >= 2 {
            for &p in &preds {
                let mut runner = p;
                while runner != idom[&b] {
                    frontiers.entry(runner).or_default().insert(b);
                    match idom.get(&runner) {
                        Some(&d) if d != runner => runner = d,
                        _ => break,
                    }
                }
            }
        }
    }

    frontiers
}

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

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

    // Classic textbook CFG:
    //   0→1, 0→2, 1→3, 2→3, 3→4
    fn diamond() -> Graph {
        let mut g = Graph::new(GraphConfig::directed());
        for _ in 0..5 { g.add_node(); }
        g.add_edge(0, 1, None);
        g.add_edge(0, 2, None);
        g.add_edge(1, 3, None);
        g.add_edge(2, 3, None);
        g.add_edge(3, 4, None);
        g
    }

    // ── immediate_dominators ──────────────────────────────────────────────────

    #[test]
    fn test_idom_chain() {
        // 0→1→2→3: each node's idom is its predecessor
        let mut g = Graph::new(GraphConfig::directed());
        for _ in 0..4 { g.add_node(); }
        for i in 0..3u64 { g.add_edge(i, i + 1, None); }
        let idom = immediate_dominators(&g, 0);
        assert_eq!(idom[&0], 0);
        assert_eq!(idom[&1], 0);
        assert_eq!(idom[&2], 1);
        assert_eq!(idom[&3], 2);
    }

    #[test]
    fn test_idom_diamond() {
        let g = diamond();
        let idom = immediate_dominators(&g, 0);
        // 0 dominates itself
        assert_eq!(idom[&0], 0);
        // 1's idom is 0, 2's idom is 0
        assert_eq!(idom[&1], 0);
        assert_eq!(idom[&2], 0);
        // 3 merges 1 and 2: idom is 0
        assert_eq!(idom[&3], 0);
        // 4's idom is 3
        assert_eq!(idom[&4], 3);
    }

    #[test]
    fn test_idom_single_node() {
        let mut g = Graph::new(GraphConfig::directed());
        g.add_node();
        let idom = immediate_dominators(&g, 0);
        assert_eq!(idom[&0], 0);
    }

    #[test]
    fn test_idom_unreachable_not_in_map() {
        // Node 2 unreachable from 0
        let mut g = Graph::new(GraphConfig::directed());
        for _ in 0..3 { g.add_node(); }
        g.add_edge(0, 1, None);
        // node 2 has no incoming edge from 0 or 1
        let idom = immediate_dominators(&g, 0);
        assert!(!idom.contains_key(&2));
    }

    // ── dominance_frontiers ───────────────────────────────────────────────────

    #[test]
    fn test_df_diamond() {
        let g = diamond();
        let idom = immediate_dominators(&g, 0);
        let df = dominance_frontiers(&g, &idom);
        // Node 3 is the join point: it should be in df[1] and df[2]
        assert!(df[&1].contains(&3), "df[1] should contain 3, got {:?}", df[&1]);
        assert!(df[&2].contains(&3), "df[2] should contain 3, got {:?}", df[&2]);
    }

    #[test]
    fn test_df_chain() {
        // Linear chain: no join points → all frontiers empty
        let mut g = Graph::new(GraphConfig::directed());
        for _ in 0..4 { g.add_node(); }
        for i in 0..3u64 { g.add_edge(i, i + 1, None); }
        let idom = immediate_dominators(&g, 0);
        let df = dominance_frontiers(&g, &idom);
        for (_, frontier) in &df {
            assert!(frontier.is_empty(), "chain has no DF, got {:?}", frontier);
        }
    }
}
