dndtree 0.2.0

//   NOTE: This test setup is from earlier commits where all variants were used and is left this
//         way for reference and continuity
//
use std::fs::File;
use std::io::BufRead;
use std::io::BufReader;
use std::path::PathBuf;

use dndtree::DNDTree;
use dndtree::bridge::ffi;
use idtree::IDTree;
use nohash_hasher::{IntMap, IntSet};
use rand::RngExt;
use rand::SeedableRng;
use rand::rngs::StdRng;

fn setup_cpp_tree(n: usize, edges: &[(usize, usize)]) -> cxx::UniquePtr<ffi::CPPDNDTree> {
    let mut adj = vec![vec![]; n];
    for &(u, v) in edges {
        if u < n && v < n {
            adj[u].push(v as i32);
            adj[v].push(u as i32);
        }
    }

    let mut degrees = Vec::with_capacity(n);
    let mut flat_neighbors = Vec::new();
    for neighbors in &adj {
        degrees.push(neighbors.len() as i32);
        for &v in neighbors {
            flat_neighbors.push(v);
        }
    }

    ffi::new_cpp_dndtree_from_flat_adj(n as i32, &degrees, &flat_neighbors, true)
}

fn make_adj_usize(n: usize, edges: &[(usize, usize)]) -> IntMap<usize, IntSet<usize>> {
    let mut adj: IntMap<usize, IntSet<usize>> = IntMap::default();
    for i in 0..n {
        adj.insert(i, IntSet::default());
    }
    for &(u, v) in edges {
        adj.get_mut(&u).unwrap().insert(v);
        adj.get_mut(&v).unwrap().insert(u);
    }
    adj
}

fn connected_idtree(tree: &mut IDTree, u: usize, v: usize) -> bool {
    tree.query(u, v)
}

fn connected_dnd(tree: &mut DNDTree, u: usize, v: usize) -> bool {
    tree.query(u, v)
}

fn make_adj(n: usize) -> IntMap<usize, IntSet<usize>> {
    let mut adj: IntMap<usize, IntSet<usize>> = IntMap::default();
    for i in 0..n {
        adj.insert(i, IntSet::default());
    }
    adj
}

fn make_edges_for_nodes(node_count: usize, edge_count: usize) -> Vec<(usize, usize)> {
    use rand::SeedableRng;
    use rand::rngs::StdRng;
    use std::collections::HashSet;

    let mut rng = StdRng::seed_from_u64(12345);
    let mut edges = Vec::with_capacity(edge_count);
    let mut seen = HashSet::with_capacity(edge_count);

    let max_possible = if node_count < 2 {
        0
    } else {
        (node_count * (node_count - 1)) / 2
    };
    let target = edge_count.min(max_possible);

    while edges.len() < target {
        let u = rng.random_range(0..node_count);
        let v = rng.random_range(0..node_count);

        if u != v {
            let canonical = if u < v { (u, v) } else { (v, u) };
            if seen.insert(canonical) {
                edges.push((u, v));
            }
        }
    }
    edges
}

fn make_caterpillar_graph(
    n: usize,
    spine_length_ratio: f64, // 0.1 = short spine, 0.5 = half nodes on spine
    extra_edges_ratio: f64,  // how many additional random chords
) -> IntMap<usize, IntSet<usize>> {
    let mut adj: IntMap<usize, IntSet<usize>> = IntMap::default();
    for i in 0..n {
        adj.insert(i, IntSet::default());
    }

    let mut rng = StdRng::seed_from_u64(42);

    // 1. Create the long spine (backbone path)
    let spine_len = (n as f64 * spine_length_ratio).max(10.0).min(n as f64) as usize;
    let mut spine = Vec::with_capacity(spine_len);
    for i in 0..spine_len {
        spine.push(i);
        if i > 0 {
            let prev = spine[i - 1];
            adj.get_mut(&prev).unwrap().insert(i);
            adj.get_mut(&i).unwrap().insert(prev);
        }
    }

    // 2. Attach remaining nodes as leaves or small trees to the spine
    let mut next_node = spine_len;
    while next_node < n {
        // Pick random spine node to attach to
        let attach_to = spine[rng.random_range(0..spine.len())];

        // Attach a small chain (1–4 nodes) to make subtrees deeper
        let chain_len = rng.random_range(1..=4);
        let mut prev = attach_to;
        for _ in 0..chain_len {
            if next_node >= n {
                break;
            }
            adj.get_mut(&prev).unwrap().insert(next_node);
            adj.get_mut(&next_node).unwrap().insert(prev);
            prev = next_node;
            next_node += 1;
        }
    }

    // 3. Add a few random chords (keep connectivity high but allow splits)
    let extra_count = (n as f64 * extra_edges_ratio) as usize;
    for _ in 0..extra_count {
        let u = rng.random_range(0..n);
        let v = rng.random_range(0..n);
        if u != v && !adj.get(&u).unwrap().contains(&v) {
            adj.get_mut(&u).unwrap().insert(v);
            adj.get_mut(&v).unwrap().insert(u);
        }
    }

    adj
}

// MARK: MTX handling
struct MtxData {
    all_edges: Vec<(usize, usize)>,
    empty_map_i32: IntMap<i32, IntSet<i32>>,
    empty_map_usize: IntMap<usize, IntSet<usize>>,
}

impl MtxData {
    fn new(filename: &str) -> Self {
        let adj_list = MtxData::load_graph(filename);
        let mut adj_dict: IntMap<i32, IntSet<i32>> = IntMap::default();
        for &u in adj_list.keys() {
            adj_dict.entry(u).or_default();
            for &v in adj_list.get(&u).unwrap() {
                adj_dict.entry(v).or_default();
                adj_dict.get_mut(&u).unwrap().insert(v);
                adj_dict.get_mut(&v).unwrap().insert(u);
            }
        }

        let mut empty_map_i32: IntMap<i32, IntSet<i32>> = IntMap::default();
        for &u in adj_list.keys() {
            empty_map_i32.entry(u).or_default();
            for &v in adj_list.get(&u).unwrap() {
                empty_map_i32.entry(v).or_default();
            }
        }

        let mut empty_map_usize: IntMap<usize, IntSet<usize>> = IntMap::default();
        for &u in adj_list.keys() {
            empty_map_usize.entry(u as usize).or_default();
            for &v in adj_list.get(&u).unwrap() {
                empty_map_usize.entry(v as usize).or_default();
            }
        }

        let all_edges: Vec<(usize, usize)> = adj_list
            .iter()
            .flat_map(|(&u, neighbors)| neighbors.iter().map(move |&v| (u as usize, v as usize)))
            .collect();

        MtxData {
            all_edges,
            empty_map_i32,
            empty_map_usize,
        }
    }

    fn load_graph(filename: &str) -> IntMap<i32, Vec<i32>> {
        let manifest_dir = env!("CARGO_MANIFEST_DIR");
        let mtx_path = PathBuf::from(manifest_dir)
            .join("benches")
            .join("data")
            .join(filename);
        let reader = BufReader::new(File::open(mtx_path).expect("MTX file missing"));

        let mut adj_list: IntMap<i32, Vec<i32>> = IntMap::default();
        let mut data_started = false;

        for line in reader.lines().map_while(Result::ok) {
            let trimmed = line.trim();
            if trimmed.is_empty() || trimmed.starts_with('%') {
                continue;
            }
            if !data_started {
                data_started = true;
                continue;
            }

            let parts: Vec<&str> = trimmed.split_whitespace().collect();
            if parts.len() >= 2 {
                let mut u: i32 = parts[0].parse().unwrap();
                let mut v: i32 = parts[1].parse().unwrap();
                // Canonicalize
                if u > v {
                    std::mem::swap(&mut u, &mut v);
                }
                // MTX is 1-based; decrement to 0-based.
                // Only add once as the graph is undirected and insert_edge handles symmetry.
                adj_list.entry(u - 1).or_default().push(v - 1);
            }
        }
        adj_list
    }
}

// MARK: Data Preparation

mod with_dsu {
    use log::debug;

    use super::*;

    #[test]
    fn test_basic_insert_delete_query() {
        let edges = vec![(0, 1), (1, 2), (2, 3)];
        let adj = make_adj_usize(4, &edges);
        let mut t = DNDTree::new(&adj);

        assert!(t.query(0, 3), "query 1");
        t.delete_edge(1, 2);
        assert!(!t.query(0, 3), "query 2");
        t.insert_edge(1, 2);
        assert!(t.query(0, 3), "query 3");
    }

    #[test]
    fn test_unlink_splits_correctly() {
        let edges = vec![(0, 1), (1, 2), (2, 3)];
        let adj = make_adj_usize(4, &edges);
        let mut t = DNDTree::new(&adj);

        t.delete_edge(1, 2);
        assert!(t.query(0, 1));
        assert!(!t.query(0, 3));
        assert!(t.query(2, 3));
    }

    #[test]
    fn test_replacement_edge_found() {
        let edges = vec![(0, 1), (1, 2), (2, 3), (0, 3)];
        let adj = make_adj_usize(4, &edges);
        let mut t = DNDTree::new(&adj);

        let r = t.delete_edge(1, 2);
        assert_eq!(r, 1);
        assert!(t.query(1, 2));
        assert!(t.query(0, 3));
    }

    #[test]
    fn test_replacement_edge_not_found() {
        let edges = vec![(0, 1), (1, 2), (2, 3)];
        let adj = make_adj_usize(4, &edges);
        let mut t = DNDTree::new(&adj);

        let r = t.delete_edge(1, 2);
        assert_eq!(r, 2);
        assert!(!t.query(0, 3));
    }

    #[test]
    fn test_dndtree_matches_idtree() {
        let mut rng = StdRng::seed_from_u64(99999);
        let n = 50;
        let mut edges = vec![];

        for _ in 0..100 {
            let u = rng.random_range(0..n);
            let v = rng.random_range(0..n);
            if u != v {
                edges.push((u, v));
            }
        }

        let adj_dnd = make_adj_usize(n, &edges);
        let adj_id = make_adj_usize(n, &edges);

        let mut dnd = DNDTree::new(&adj_dnd);
        let mut idt = IDTree::new(&adj_id);

        for _ in 0..200 {
            let u = rng.random_range(0..n);
            let v = rng.random_range(0..n);

            let op = rng.random_range(0..3);
            match op {
                0 => {
                    dnd.insert_edge(u, v);
                    idt.insert_edge(u, v);
                }
                1 => {
                    dnd.delete_edge(u, v);
                    idt.delete_edge(u, v);
                }
                _ => {}
            }

            for _ in 0..20 {
                let a = rng.random_range(0..n);
                let b = rng.random_range(0..n);
                assert_eq!(
                    connected_dnd(&mut dnd, a, b),
                    connected_idtree(&mut idt, a, b)
                );
            }
        }
    }

    #[test]
    fn test_mixed_ops_query_heavy() {
        let n = 1_000;
        let query_factor = 10;

        let edges = make_edges_for_nodes(n, n * 2);
        let (present_edges, absent_edges) = edges.split_at(n);

        let mut adj = make_adj(n);
        for &(u, v) in present_edges.iter() {
            adj.get_mut(&(u)).unwrap().insert(v);
            adj.get_mut(&(v)).unwrap().insert(u);
        }

        let mut tree = DNDTree::new(&adj);

        let mut present: Vec<usize> = (0..n).collect();
        let mut absent: Vec<usize> = (0..n).collect();

        for i in 0..n {
            let pi = present[i];
            let (du, dv) = present_edges[pi];
            tree.delete_edge(du, dv);

            for q in 0..query_factor {
                let qi = present[(i + q) % n];
                let (qu, qv) = present_edges[qi];
                let _ = tree.query(qu, qv);
            }

            let ai = absent[i];
            let (iu, iv) = absent_edges[ai];
            tree.insert_edge(iu, iv);

            present[i] = ai;
            absent[i] = pi;
        }
    }

    #[test]
    fn mixed_ops_query_heavy_catgraph() {
        const QUERY_FACTOR: f64 = 0.05;

        use rand::SeedableRng;
        use rand::rngs::StdRng;

        let n = 20_000;
        let edges = make_edges_for_nodes(n, n * 2);
        let mut rng = StdRng::seed_from_u64(12345);

        let (present_edges, absent_edges) = edges.split_at(n);

        let adj = make_caterpillar_graph(n, 0.3, 0.05); // spine ~30% of nodes, 5% extra chords

        // Pre-select random endpoint pairs (not edges)
        let num_query_pairs = (QUERY_FACTOR * n as f64) as usize;
        let mut query_pairs = Vec::with_capacity(num_query_pairs);
        for _ in 0..num_query_pairs {
            let qu = rng.random_range(0..n);
            let qv = rng.random_range(0..n);
            query_pairs.push((qu, qv));
        }

        let mut tree = DNDTree::new(&adj);

        let mut present: Vec<usize> = (0..n).collect();
        let mut absent: Vec<usize> = (0..n).collect();

        for i in 0..n {
            let pi = present[i];
            let (du, dv) = present_edges[pi];
            tree.delete_edge(du, dv);

            for &(qu, qv) in &query_pairs {
                let _ = tree.query(qu, qv);
            }

            let ai = absent[i];
            let (iu, iv) = absent_edges[ai];
            tree.insert_edge(iu, iv);

            present[i] = ai;
            absent[i] = pi;
        }
    }

    // This test will fail if sorting of neighbors is not applied
    #[test]
    #[ignore]
    fn test_dndtree_matches_idtree_mtx() {
        // let _ = simple_logger::SimpleLogger::new()
        //     .with_level(log::LevelFilter::Debug)
        //     .init();

        let mtx_data = MtxData::new("bdo_exploration_graph.mtx");
        let all_nodes = mtx_data.empty_map_i32.keys().collect::<Vec<_>>();

        let mut dnd = DNDTree::new(&mtx_data.empty_map_usize);
        let mut idt = IDTree::new(&mtx_data.empty_map_usize);

        for &(u, v) in mtx_data.all_edges.iter() {
            debug!(
                "----------------------------------------------\n inserting edge u = {u}, v = {v}\n----------------------------------------------"
            );
            debug!("-----DNDTree no dsu no compression");
            let dnd_res = dnd.insert_edge(u, v);
            debug!("\n-----IDTree");
            let idt_res = idt.insert_edge(u, v);

            // Validate all nodes state
            let mut divergent_state = false;
            for &u in all_nodes.iter() {
                let dnd_node_data = dnd.get_node_data(*u as usize);
                let idt_node_data = idt.get_node_data(*u as usize);
                if dnd_node_data.parent != idt_node_data.parent as usize
                    || dnd_node_data.subtree_size != idt_node_data.subtree_size
                    || dnd_node_data.neighbors.len() != idt_node_data.neighbors.len()
                {
                    divergent_state = true;
                    debug!(
                        "post insert state of u: {}\n  dnd_node_data: {:?}\n  idt_node_data: {:?}",
                        u, dnd_node_data.neighbors, idt_node_data.neighbors
                    );
                }
            }
            // for &i in all_nodes.iter() {
            //     for &j in all_nodes.iter() {
            //         let dnd_conn = dnd.query(*i as usize, *j as usize);
            //         let idt_conn = idt.query(*i as usize, *j as usize);
            //         if dnd_conn != idt_conn {
            //             println!(
            //                 "FIRST connectivity mismatch after op on ({u}, {v}) at pair ({}, {}): dnd={}, idt={}",
            //                 i, j, dnd_conn, idt_conn
            //             );
            //             println!("Connectivity diverged");
            //             divergent_state = true;
            //         }
            //     }
            // }

            if divergent_state {
                panic!("DNDTree and IDTree diverged!");
            }
            assert!(
                dnd_res == idt_res,
                "Insert results don't match u = {u}, v = {v}, dnd_res = {dnd_res}, idt_res = {idt_res}"
            );
        }
        for &(u, v) in mtx_data.all_edges.iter() {
            debug!(
                "----------------------------------------------\n deleting edge u = {u}, v = {v}\n----------------------------------------------"
            );
            debug!("-----DNDTree dsu no compression");
            let dnd_res = dnd.delete_edge(u, v);
            debug!("\n-----IDTree");
            let idt_res = idt.delete_edge(u, v);

            // Validate all nodes state
            let mut divergent_state = false;
            for &u in all_nodes.iter() {
                let dnd_node_data = dnd.get_node_data(*u as usize);
                let idt_node_data = idt.get_node_data(*u as usize);
                if dnd_node_data.parent != idt_node_data.parent as usize
                    || dnd_node_data.subtree_size != idt_node_data.subtree_size
                    || dnd_node_data.neighbors.len() != idt_node_data.neighbors.len()
                {
                    divergent_state = true;
                    debug!(
                        "post insert state of u: {}\n  dnd_node_data: {:?}\n  idt_node_data: {:?}",
                        u, dnd_node_data, idt_node_data
                    );
                }
            }
            // for &i in all_nodes.iter() {
            //     for &j in all_nodes.iter() {
            //         let dnd_conn = dnd.query(*i as usize, *j as usize);
            //         let idt_conn = idt.query(*i as usize, *j as usize);
            //         if dnd_conn != idt_conn {
            //             println!(
            //                 "FIRST connectivity mismatch after op on ({u}, {v}) at pair ({}, {}): dnd={}, idt={}",
            //                 i, j, dnd_conn, idt_conn
            //             );
            //             println!("Connectivity diverged");
            //             divergent_state = true;
            //         }
            //     }
            // }

            if divergent_state {
                panic!("DNDTree and IDTree diverged!");
            }
            assert!(
                dnd_res == idt_res,
                "Delete results don't match u = {u}, v = {v}, dnd_res = {dnd_res}, idt_res = {idt_res}"
            );
            let dnd_res = dnd.query(u, v);
            let idt_res = idt.query(u, v);
            assert!(
                dnd_res == idt_res,
                "Query results don't match after delete u = {u}, v = {v}, dnd_res = {dnd_res}, idt_res = {idt_res}"
            );
        }

        for &(u, v) in mtx_data.all_edges.iter() {
            assert_eq!(
                connected_dnd(&mut dnd, u, v),
                connected_idtree(&mut idt, u, v)
            );
        }
    }
}

mod cpp_tests {
    use super::*;
    use cxx::UniquePtr;
    use dndtree::bridge::ffi;
    use log::debug;

    fn verify_idtree_topology_sync(cpp: &UniquePtr<ffi::CPPDNDTree>, idt: &IDTree, n: usize) {
        for i in 0..n {
            let cpp_p = cpp.get_tree_parent(i as i32);
            let idt_p = idt.get_parent(i); // Assuming this is your IDTree getter

            if cpp_p != idt_p as i32 {
                panic!(
                    "TOPOLOGY DRIFT at Node {}: CPP parent = {}, IDT parent = {}",
                    i, cpp_p, idt_p
                );
            }
        }
    }

    // This test will fail if sorting of neighbors is not applied to insert/delete neighbors.
    #[test]
    #[ignore]
    fn test_dndtree_matches_idtree_mtx() {
        use std::io::Write; // Required for flushing

        // let _ = simple_logger::SimpleLogger::new()
        //     .with_level(log::LevelFilter::Debug)
        //     .init();

        // dndtree::bridge::ffi::set_cpp_trace(true);

        let mtx_data = MtxData::new("bdo_exploration_graph.mtx");
        let all_nodes = mtx_data.empty_map_i32.keys().collect::<Vec<_>>();
        let node_count = all_nodes.len();

        let cpp = setup_cpp_tree(all_nodes.len(), &[]);
        let mut idt = IDTree::new(&mtx_data.empty_map_usize);

        // Insertion Phase
        for &(u, v) in mtx_data.all_edges.iter() {
            debug!(
                "----------------------------------------------\n inserting edge u = {u}, v = {v}\n----------------------------------------------"
            );
            let _ = std::io::stdout().flush();

            debug!("-----CPPDNDTree");
            let _ = std::io::stdout().flush();

            let cpp_res = cpp.insert_edge(u as i32, v as i32);

            debug!("\n-----IDTree");
            let _ = std::io::stdout().flush();

            let idt_res = idt.insert_edge(u as usize, v as usize);

            verify_idtree_topology_sync(&cpp, &idt, node_count);

            if cpp_res != idt_res as i32 {
                panic!(
                    "Insert results don't match u = {u}, v = {v}, cpp_res = {cpp_res}, idt_res = {idt_res}"
                );
            }
        }

        // Deletion Phase with Connectivity Checks
        for &(u, v) in mtx_data.all_edges.iter() {
            debug!(
                "----------------------------------------------\n deleting edge u = {u}, v = {v}\n----------------------------------------------"
            );
            debug!("-----CPPDNDTree");
            let _ = std::io::stdout().flush();

            let cpp_res = cpp.delete_edge(u as i32, v as i32);

            debug!("\n-----IDTree");
            let _ = std::io::stdout().flush();

            // Perform the operation that is likely causing the trace cutoff
            let idt_res = idt.delete_edge(u as usize, v as usize);

            verify_idtree_topology_sync(&cpp, &idt, node_count);

            // Critical flush: If idt.delete_edge panicked or corrupted the state,
            // we want the internal logger to have pushed its buffer out.
            let _ = std::io::stdout().flush();

            if cpp_res != idt_res as i32 {
                debug!(
                    "Delete results don't match u = {u}, v = {v}, cpp_res = {cpp_res}, idt_res = {idt_res}"
                );
                let _ = std::io::stdout().flush();
                panic!("CPPDNDTree and IDTree diverged!");
            }

            let cpp_conn = cpp.query(u as i32, v as i32);
            let idt_conn = idt.query(u as usize, v as usize);
            if cpp_conn != idt_conn {
                debug!(
                    "Query results don't match after delete u = {u}, v = {v}, cpp_res = {cpp_conn}, idt_res = {idt_conn}"
                );
                let _ = std::io::stdout().flush();
                panic!("CPPDNDTree and IDTree diverged!");
            }
        }

        // Final Connectivity Sweep
        for &(u, v) in mtx_data.all_edges.iter() {
            assert_eq!(
                cpp.query(u as i32, v as i32),
                idt.query(u as usize, v as usize),
                "Final connectivity mismatch for edge ({}, {})",
                u,
                v
            );
        }
    }
}