use alloc::{
    collections::{BTreeMap, BTreeSet},
    vec::Vec,
};

use assert_matches::assert_matches;
use proptest::prelude::*;
use rand::{Rng, prelude::IteratorRandom, rng};

use super::{
    COLS_PER_SUBTREE, InnerNode, NodeIndex, NodeMutations, PairComputations, RpoDigest, SMT_DEPTH,
    SUBTREE_DEPTH, Smt, SmtLeaf, SparseMerkleTree, SubtreeLeaf, SubtreeLeavesIter, UnorderedMap,
    build_subtree,
};
use crate::{
    EMPTY_WORD, ONE, Word, ZERO,
    merkle::{LeafIndex, MerkleError, smt::Felt},
};

fn smtleaf_to_subtree_leaf(leaf: &SmtLeaf) -> SubtreeLeaf {
    SubtreeLeaf {
        col: leaf.index().index.value(),
        hash: leaf.hash(),
    }
}

#[test]
fn test_sorted_pairs_to_leaves() {
    let entries: Vec<(RpoDigest, Word)> = vec![
        // Subtree 0.
        (RpoDigest::new([ONE, ONE, ONE, Felt::new(16)]), [ONE; 4]),
        (RpoDigest::new([ONE, ONE, ONE, Felt::new(17)]), [ONE; 4]),
        // Leaf index collision.
        (RpoDigest::new([ONE, ONE, Felt::new(10), Felt::new(20)]), [ONE; 4]),
        (RpoDigest::new([ONE, ONE, Felt::new(20), Felt::new(20)]), [ONE; 4]),
        // Subtree 1. Normal single leaf again.
        (RpoDigest::new([ONE, ONE, ONE, Felt::new(400)]), [ONE; 4]), // Subtree boundary.
        (RpoDigest::new([ONE, ONE, ONE, Felt::new(401)]), [ONE; 4]),
        // Subtree 2. Another normal leaf.
        (RpoDigest::new([ONE, ONE, ONE, Felt::new(1024)]), [ONE; 4]),
    ];

    let control = Smt::with_entries_sequential(entries.clone()).unwrap();
    let control_leaves: Vec<SmtLeaf> = {
        let mut entries_iter = entries.iter().cloned();
        let mut next_entry = || entries_iter.next().unwrap();
        let control_leaves = vec![
            // Subtree 0.
            SmtLeaf::Single(next_entry()),
            SmtLeaf::Single(next_entry()),
            SmtLeaf::new_multiple(vec![next_entry(), next_entry()]).unwrap(),
            // Subtree 1.
            SmtLeaf::Single(next_entry()),
            SmtLeaf::Single(next_entry()),
            // Subtree 2.
            SmtLeaf::Single(next_entry()),
        ];
        assert_eq!(entries_iter.next(), None);
        control_leaves
    };

    let control_subtree_leaves: Vec<Vec<SubtreeLeaf>> = {
        let mut control_leaves_iter = control_leaves.iter();
        let mut next_leaf = || control_leaves_iter.next().unwrap();
        let control_subtree_leaves: Vec<Vec<SubtreeLeaf>> = [
            // Subtree 0.
            vec![next_leaf(), next_leaf(), next_leaf()],
            // Subtree 1.
            vec![next_leaf(), next_leaf()],
            // Subtree 2.
            vec![next_leaf()],
        ]
        .map(|subtree| subtree.into_iter().map(smtleaf_to_subtree_leaf).collect())
        .to_vec();
        assert_eq!(control_leaves_iter.next(), None);
        control_subtree_leaves
    };

    let subtrees: PairComputations<u64, SmtLeaf> = Smt::sorted_pairs_to_leaves(entries).unwrap();
    // This will check that the hashes, columns, and subtree assignments all match.
    assert_eq!(subtrees.leaves, control_subtree_leaves);
    // Flattening and re-separating out the leaves into subtrees should have the same result.
    let mut all_leaves: Vec<SubtreeLeaf> = subtrees.leaves.clone().into_iter().flatten().collect();
    let re_grouped: Vec<Vec<_>> = SubtreeLeavesIter::from_leaves(&mut all_leaves).collect();
    assert_eq!(subtrees.leaves, re_grouped);
    // Then finally we might as well check the computed leaf nodes too.
    let control_leaves: BTreeMap<u64, SmtLeaf> = control
        .leaves()
        .map(|(index, value)| (index.index.value(), value.clone()))
        .collect();

    for (column, test_leaf) in subtrees.nodes {
        if test_leaf.is_empty() {
            continue;
        }
        let control_leaf = control_leaves
            .get(&column)
            .unwrap_or_else(|| panic!("no leaf node found for column {column}"));
        assert_eq!(control_leaf, &test_leaf);
    }
}

// Helper for the below tests.
fn generate_entries(pair_count: u64) -> Vec<(RpoDigest, Word)> {
    (0..pair_count)
        .map(|i| {
            let leaf_index = ((i as f64 / pair_count as f64) * (pair_count as f64)) as u64;
            let key = RpoDigest::new([ONE, ONE, Felt::new(i), Felt::new(leaf_index)]);
            let value = [ONE, ONE, ONE, Felt::new(i)];
            (key, value)
        })
        .collect()
}

fn generate_updates(entries: Vec<(RpoDigest, Word)>, updates: usize) -> Vec<(RpoDigest, Word)> {
    const REMOVAL_PROBABILITY: f64 = 0.2;
    let mut rng = rng();
    // Assertion to ensure input keys are unique
    assert!(
        entries.iter().map(|(key, _)| key).collect::<BTreeSet<_>>().len() == entries.len(),
        "Input entries contain duplicate keys!"
    );
    let mut sorted_entries: Vec<(RpoDigest, Word)> = entries
        .into_iter()
        .choose_multiple(&mut rng, updates)
        .into_iter()
        .map(|(key, _)| {
            let value = if rng.random_bool(REMOVAL_PROBABILITY) {
                EMPTY_WORD
            } else {
                [ONE, ONE, ONE, Felt::new(rng.random())]
            };
            (key, value)
        })
        .collect();
    sorted_entries.sort_by_key(|(key, _)| Smt::key_to_leaf_index(key).value());
    sorted_entries
}

#[test]
fn test_single_subtree() {
    // A single subtree's worth of leaves.
    const PAIR_COUNT: u64 = COLS_PER_SUBTREE;
    let entries = generate_entries(PAIR_COUNT);
    let control = Smt::with_entries_sequential(entries.clone()).unwrap();
    // `entries` should already be sorted by nature of how we constructed it.
    let leaves = Smt::sorted_pairs_to_leaves(entries).unwrap().leaves;
    let leaves = leaves.into_iter().next().unwrap();
    let (first_subtree, subtree_root) = build_subtree(leaves, SMT_DEPTH, SMT_DEPTH);
    assert!(!first_subtree.is_empty());
    // The inner nodes computed from that subtree should match the nodes in our control tree.
    for (index, node) in first_subtree.into_iter() {
        let control = control.get_inner_node(index);
        assert_eq!(
            control, node,
            "subtree-computed node at index {index:?} does not match control",
        );
    }
    // The root returned should also match the equivalent node in the control tree.
    let control_root_index =
        NodeIndex::new(SMT_DEPTH - SUBTREE_DEPTH, subtree_root.col).expect("Valid root index");
    let control_root_node = control.get_inner_node(control_root_index);
    let control_hash = control_root_node.hash();
    assert_eq!(
        control_hash, subtree_root.hash,
        "Subtree-computed root at index {control_root_index:?} does not match control"
    );
}

// Test that not just can we compute a subtree correctly, but we can feed the results of one
// subtree into computing another. In other words, test that `build_subtree()` is correctly
// composable.
#[test]
fn test_two_subtrees() {
    // Two subtrees' worth of leaves.
    const PAIR_COUNT: u64 = COLS_PER_SUBTREE * 2;
    let entries = generate_entries(PAIR_COUNT);
    let control = Smt::with_entries_sequential(entries.clone()).unwrap();
    let PairComputations { leaves, .. } = Smt::sorted_pairs_to_leaves(entries).unwrap();
    // With two subtrees' worth of leaves, we should have exactly two subtrees.
    let [first, second]: [Vec<_>; 2] = leaves.try_into().unwrap();
    assert_eq!(first.len() as u64, PAIR_COUNT / 2);
    assert_eq!(first.len(), second.len());
    let mut current_depth = SMT_DEPTH;
    let mut next_leaves: Vec<SubtreeLeaf> = Default::default();
    let (first_nodes, first_root) = build_subtree(first, SMT_DEPTH, current_depth);
    next_leaves.push(first_root);
    let (second_nodes, second_root) = build_subtree(second, SMT_DEPTH, current_depth);
    next_leaves.push(second_root);
    // All new inner nodes + the new subtree-leaves should be 512, for one depth-cycle.
    let total_computed = first_nodes.len() + second_nodes.len() + next_leaves.len();
    assert_eq!(total_computed as u64, PAIR_COUNT);
    // Verify the computed nodes of both subtrees.
    let computed_nodes = first_nodes.clone().into_iter().chain(second_nodes);
    for (index, test_node) in computed_nodes {
        let control_node = control.get_inner_node(index);
        assert_eq!(
            control_node, test_node,
            "subtree-computed node at index {index:?} does not match control",
        );
    }
    current_depth -= SUBTREE_DEPTH;
    let (nodes, root_leaf) = build_subtree(next_leaves, SMT_DEPTH, current_depth);
    assert_eq!(nodes.len(), SUBTREE_DEPTH as usize);
    assert_eq!(root_leaf.col, 0);
    for (index, test_node) in nodes {
        let control_node = control.get_inner_node(index);
        assert_eq!(
            control_node, test_node,
            "subtree-computed node at index {index:?} does not match control",
        );
    }
    let index = NodeIndex::new(current_depth - SUBTREE_DEPTH, root_leaf.col).unwrap();
    let control_root = control.get_inner_node(index).hash();
    assert_eq!(control_root, root_leaf.hash, "Root mismatch");
}

#[test]
fn test_singlethreaded_subtrees() {
    const PAIR_COUNT: u64 = COLS_PER_SUBTREE * 64;
    let entries = generate_entries(PAIR_COUNT);
    let control = Smt::with_entries_sequential(entries.clone()).unwrap();
    let mut accumulated_nodes: BTreeMap<NodeIndex, InnerNode> = Default::default();
    let PairComputations {
        leaves: mut leaf_subtrees,
        nodes: test_leaves,
    } = Smt::sorted_pairs_to_leaves(entries).unwrap();
    for current_depth in (SUBTREE_DEPTH..=SMT_DEPTH).step_by(SUBTREE_DEPTH as usize).rev() {
        // There's no flat_map_unzip(), so this is the best we can do.
        let (nodes, mut subtree_roots): (Vec<UnorderedMap<_, _>>, Vec<SubtreeLeaf>) = leaf_subtrees
            .into_iter()
            .enumerate()
            .map(|(i, subtree)| {
                // Pre-assertions.
                assert!(
                    subtree.is_sorted(),
                    "subtree {i} at bottom-depth {current_depth} is not sorted",
                );
                assert!(
                    !subtree.is_empty(),
                    "subtree {i} at bottom-depth {current_depth} is empty!",
                );
                // Do actual things.
                let (nodes, subtree_root) = build_subtree(subtree, SMT_DEPTH, current_depth);
                // Post-assertions.
                for (&index, test_node) in nodes.iter() {
                    let control_node = control.get_inner_node(index);
                    assert_eq!(
                        test_node, &control_node,
                        "depth {} subtree {}: test node does not match control at index {:?}",
                        current_depth, i, index,
                    );
                }
                (nodes, subtree_root)
            })
            .unzip();
        // Update state between each depth iteration.
        leaf_subtrees = SubtreeLeavesIter::from_leaves(&mut subtree_roots).collect();
        accumulated_nodes.extend(nodes.into_iter().flatten());
        assert!(!leaf_subtrees.is_empty(), "on depth {current_depth}");
    }
    // Make sure the true leaves match, first checking length and then checking each individual
    // leaf.
    let control_leaves: BTreeMap<_, _> = control.leaves().collect();
    let control_leaves_len = control_leaves.len();
    let test_leaves_len = test_leaves.len();
    assert_eq!(test_leaves_len, control_leaves_len);
    for (col, ref test_leaf) in test_leaves {
        let index = LeafIndex::new_max_depth(col);
        let &control_leaf = control_leaves.get(&index).unwrap();
        assert_eq!(test_leaf, control_leaf, "test leaf at column {col} does not match control");
    }
    // Make sure the inner nodes match, checking length first and then each individual leaf.
    let control_nodes_len = control.inner_nodes().count();
    let test_nodes_len = accumulated_nodes.len();
    assert_eq!(test_nodes_len, control_nodes_len);
    for (index, test_node) in accumulated_nodes.clone() {
        let control_node = control.get_inner_node(index);
        assert_eq!(test_node, control_node, "test node does not match control at {index:?}");
    }
    // After the last iteration of the above for loop, we should have the new root node actually
    // in two places: one in `accumulated_nodes`, and the other as the "next leaves" return from
    // `build_subtree()`. So let's check both!
    let control_root = control.get_inner_node(NodeIndex::root());
    // That for loop should have left us with only one leaf subtree...
    let [leaf_subtree]: [Vec<_>; 1] = leaf_subtrees.try_into().unwrap();
    // which itself contains only one 'leaf'...
    let [root_leaf]: [SubtreeLeaf; 1] = leaf_subtree.try_into().unwrap();
    // which matches the expected root.
    assert_eq!(control.root(), root_leaf.hash);
    // Likewise `accumulated_nodes` should contain a node at the root index...
    assert!(accumulated_nodes.contains_key(&NodeIndex::root()));
    // and it should match our actual root.
    let test_root = accumulated_nodes.get(&NodeIndex::root()).unwrap();
    assert_eq!(control_root, *test_root);
    // And of course the root we got from each place should match.
    assert_eq!(control.root(), root_leaf.hash);
}

/// The parallel version of `test_singlethreaded_subtree()`.
#[test]
fn test_multithreaded_subtrees() {
    use rayon::prelude::*;
    const PAIR_COUNT: u64 = COLS_PER_SUBTREE * 64;
    let entries = generate_entries(PAIR_COUNT);
    let control = Smt::with_entries_sequential(entries.clone()).unwrap();
    let mut accumulated_nodes: BTreeMap<NodeIndex, InnerNode> = Default::default();
    let PairComputations {
        leaves: mut leaf_subtrees,
        nodes: test_leaves,
    } = Smt::sorted_pairs_to_leaves(entries).unwrap();
    for current_depth in (SUBTREE_DEPTH..=SMT_DEPTH).step_by(SUBTREE_DEPTH as usize).rev() {
        let (nodes, mut subtree_roots): (Vec<UnorderedMap<_, _>>, Vec<SubtreeLeaf>) = leaf_subtrees
            .into_par_iter()
            .enumerate()
            .map(|(i, subtree)| {
                // Pre-assertions.
                assert!(
                    subtree.is_sorted(),
                    "subtree {i} at bottom-depth {current_depth} is not sorted",
                );
                assert!(
                    !subtree.is_empty(),
                    "subtree {i} at bottom-depth {current_depth} is empty!",
                );
                let (nodes, subtree_root) = build_subtree(subtree, SMT_DEPTH, current_depth);
                // Post-assertions.
                for (&index, test_node) in nodes.iter() {
                    let control_node = control.get_inner_node(index);
                    assert_eq!(
                        test_node, &control_node,
                        "depth {} subtree {}: test node does not match control at index {:?}",
                        current_depth, i, index,
                    );
                }
                (nodes, subtree_root)
            })
            .unzip();
        leaf_subtrees = SubtreeLeavesIter::from_leaves(&mut subtree_roots).collect();
        accumulated_nodes.extend(nodes.into_iter().flatten());
        assert!(!leaf_subtrees.is_empty(), "on depth {current_depth}");
    }
    // Make sure the true leaves match, checking length first and then each individual leaf.
    let control_leaves: BTreeMap<_, _> = control.leaves().collect();
    let control_leaves_len = control_leaves.len();
    let test_leaves_len = test_leaves.len();
    assert_eq!(test_leaves_len, control_leaves_len);
    for (col, ref test_leaf) in test_leaves {
        let index = LeafIndex::new_max_depth(col);
        let &control_leaf = control_leaves.get(&index).unwrap();
        assert_eq!(test_leaf, control_leaf);
    }
    // Make sure the inner nodes match, checking length first and then each individual leaf.
    let control_nodes_len = control.inner_nodes().count();
    let test_nodes_len = accumulated_nodes.len();
    assert_eq!(test_nodes_len, control_nodes_len);
    for (index, test_node) in accumulated_nodes.clone() {
        let control_node = control.get_inner_node(index);
        assert_eq!(test_node, control_node, "test node does not match control at {index:?}");
    }
    // After the last iteration of the above for loop, we should have the new root node actually
    // in two places: one in `accumulated_nodes`, and the other as the "next leaves" return from
    // `build_subtree()`. So let's check both!
    let control_root = control.get_inner_node(NodeIndex::root());
    // That for loop should have left us with only one leaf subtree...
    let [leaf_subtree]: [_; 1] = leaf_subtrees.try_into().unwrap();
    // which itself contains only one 'leaf'...
    let [root_leaf]: [_; 1] = leaf_subtree.try_into().unwrap();
    // which matches the expected root.
    assert_eq!(control.root(), root_leaf.hash);
    // Likewise `accumulated_nodes` should contain a node at the root index...
    assert!(accumulated_nodes.contains_key(&NodeIndex::root()));
    // and it should match our actual root.
    let test_root = accumulated_nodes.get(&NodeIndex::root()).unwrap();
    assert_eq!(control_root, *test_root);
    // And of course the root we got from each place should match.
    assert_eq!(control.root(), root_leaf.hash);
}

#[test]
fn test_with_entries_concurrent() {
    const PAIR_COUNT: u64 = COLS_PER_SUBTREE * 64;
    let mut entries = generate_entries(PAIR_COUNT);
    let mut rng = rand::rng();

    // Set 10% of the entries to have empty words as their values.
    for _ in 0..PAIR_COUNT / 10 {
        let random_index = rng.random_range(0..PAIR_COUNT);
        entries[random_index as usize].1 = EMPTY_WORD;
    }

    let control = Smt::with_entries_sequential(entries.clone()).unwrap();
    let smt = Smt::with_entries(entries.clone()).unwrap();
    assert_eq!(smt.root(), control.root());
    assert_eq!(smt, control);
}

/// Concurrent mutations
#[test]
fn test_singlethreaded_subtree_mutations() {
    const PAIR_COUNT: u64 = COLS_PER_SUBTREE * 64;
    let entries = generate_entries(PAIR_COUNT);
    let updates = generate_updates(entries.clone(), 1000);
    let tree = Smt::with_entries_sequential(entries.clone()).unwrap();
    let control = tree.compute_mutations_sequential(updates.clone());
    let mut node_mutations = NodeMutations::default();
    let (mut subtree_leaves, new_pairs) = tree.sorted_pairs_to_mutated_subtree_leaves(updates);
    for current_depth in (SUBTREE_DEPTH..=SMT_DEPTH).step_by(SUBTREE_DEPTH as usize).rev() {
        // There's no flat_map_unzip(), so this is the best we can do.
        let (mutations_per_subtree, mut subtree_roots): (Vec<_>, Vec<_>) = subtree_leaves
            .into_iter()
            .enumerate()
            .map(|(i, subtree)| {
                // Pre-assertions.
                assert!(
                    subtree.is_sorted(),
                    "subtree {i} at bottom-depth {current_depth} is not sorted",
                );
                assert!(
                    !subtree.is_empty(),
                    "subtree {i} at bottom-depth {current_depth} is empty!",
                );
                // Calculate the mutations for this subtree.
                let (mutations_per_subtree, subtree_root) =
                    tree.build_subtree_mutations(subtree, SMT_DEPTH, current_depth);
                // Check that the mutations match the control tree.
                for (&index, mutation) in mutations_per_subtree.iter() {
                    let control_mutation = control.node_mutations().get(&index).unwrap();
                    assert_eq!(
                        control_mutation, mutation,
                        "depth {} subtree {}: mutation does not match control at index {:?}",
                        current_depth, i, index,
                    );
                }
                (mutations_per_subtree, subtree_root)
            })
            .unzip();
        subtree_leaves = SubtreeLeavesIter::from_leaves(&mut subtree_roots).collect();
        node_mutations.extend(mutations_per_subtree.into_iter().flatten());
        assert!(!subtree_leaves.is_empty(), "on depth {current_depth}");
    }
    let [subtree]: [Vec<_>; 1] = subtree_leaves.try_into().unwrap();
    let [root_leaf]: [SubtreeLeaf; 1] = subtree.try_into().unwrap();
    // Check that the new root matches the control.
    assert_eq!(control.new_root, root_leaf.hash);
    // Check that the node mutations match the control.
    assert_eq!(control.node_mutations().len(), node_mutations.len());
    for (&index, mutation) in control.node_mutations().iter() {
        let test_mutation = node_mutations.get(&index).unwrap();
        assert_eq!(test_mutation, mutation);
    }
    // Check that the new pairs match the control
    assert_eq!(control.new_pairs.len(), new_pairs.len());
    for (&key, &value) in control.new_pairs.iter() {
        let test_value = new_pairs.get(&key).unwrap();
        assert_eq!(test_value, &value);
    }
}

#[test]
fn test_compute_mutations_parallel() {
    const PAIR_COUNT: u64 = COLS_PER_SUBTREE * 64;
    let entries = generate_entries(PAIR_COUNT);
    let tree = Smt::with_entries(entries.clone()).unwrap();
    let updates = generate_updates(entries, 1000);
    let control = tree.compute_mutations_sequential(updates.clone());
    let mutations = tree.compute_mutations(updates);
    assert_eq!(mutations.root(), control.root());
    assert_eq!(mutations.old_root(), control.old_root());
    assert_eq!(mutations.node_mutations(), control.node_mutations());
    assert_eq!(mutations.new_pairs(), control.new_pairs());
}

#[test]
fn test_smt_construction_with_entries_unsorted() {
    let entries = [
        (RpoDigest::new([ONE, ONE, Felt::new(2_u64), ONE]), [ONE; 4]),
        (RpoDigest::new([ONE; 4]), [ONE; 4]),
    ];
    let control = Smt::with_entries_sequential(entries).unwrap();
    let smt = Smt::with_entries(entries).unwrap();
    assert_eq!(smt.root(), control.root());
    assert_eq!(smt, control);
}

#[test]
fn test_smt_construction_with_entries_duplicate_keys() {
    let entries = [
        (RpoDigest::new([ONE, ONE, ONE, Felt::new(16)]), [ONE; 4]),
        (RpoDigest::new([ONE; 4]), [ONE; 4]),
        (RpoDigest::new([ONE, ONE, ONE, Felt::new(16)]), [ONE; 4]),
    ];
    let expected_col = Smt::key_to_leaf_index(&entries[0].0).index.value();
    let err = Smt::with_entries(entries).unwrap_err();
    assert_matches!(err, MerkleError::DuplicateValuesForIndex(col) if col == expected_col);
}

#[test]
fn test_smt_construction_with_some_empty_values() {
    let entries = [
        (RpoDigest::new([ONE, ONE, ONE, ONE]), Smt::EMPTY_VALUE),
        (RpoDigest::new([ONE, ONE, ONE, Felt::new(2)]), [ONE; 4]),
    ];

    let result = Smt::with_entries(entries);
    assert!(result.is_ok(), "SMT construction failed with mixed empty values");

    let smt = result.unwrap();
    let control = Smt::with_entries_sequential(entries).unwrap();

    assert_eq!(smt.num_leaves(), 1);
    assert_eq!(smt.root(), control.root(), "Root hashes do not match");
    assert_eq!(smt, control, "SMTs are not equal");
}

#[test]
fn test_smt_construction_with_all_empty_values() {
    let entries = [(RpoDigest::new([ONE, ONE, ONE, ONE]), Smt::EMPTY_VALUE)];

    let result = Smt::with_entries(entries);
    assert!(result.is_ok(), "SMT construction failed with all empty values");

    let smt = result.unwrap();

    assert_eq!(
        smt.root(),
        Smt::default().root(),
        "SMT with all empty values should have the same root as the default SMT"
    );
    assert_eq!(smt, Smt::default(), "SMT with all empty values should be empty");
}

#[test]
fn test_smt_construction_with_no_entries() {
    let entries: [(RpoDigest, Word); 0] = [];

    let result = Smt::with_entries(entries);
    assert!(result.is_ok(), "SMT construction failed with no entries");

    let smt = result.unwrap();
    assert_eq!(smt, Smt::default(), "SMT with no entries should be empty");
}

fn arb_felt() -> impl Strategy<Value = Felt> {
    prop_oneof![any::<u64>().prop_map(Felt::new), Just(ZERO), Just(ONE),]
}

/// Generate entries that are guaranteed to be in different subtrees
fn generate_cross_subtree_entries() -> impl Strategy<Value = Vec<(RpoDigest, Word)>> {
    let subtree_offsets = prop::collection::vec(0..(COLS_PER_SUBTREE * 4), 1..100);

    subtree_offsets.prop_map(|offsets| {
        offsets
            .into_iter()
            .map(|base_col| {
                let key = RpoDigest::new([ONE, ONE, ONE, Felt::new(base_col)]);
                let value = [ONE, ONE, ONE, Felt::new(base_col)];
                (key, value)
            })
            .collect()
    })
}

fn arb_entries() -> impl Strategy<Value = Vec<(RpoDigest, Word)>> {
    // Combine random entries with guaranteed cross-subtree entries
    prop_oneof![
        // Original random entry generation
        prop::collection::vec(
            prop_oneof![
                // Random values case
                (
                    prop::array::uniform4(arb_felt()).prop_map(RpoDigest::new),
                    prop::array::uniform4(arb_felt())
                ),
                // Edge case values
                (
                    Just(RpoDigest::new([ONE, ONE, ONE, Felt::new(0)])),
                    Just([ONE, ONE, ONE, Felt::new(u64::MAX)])
                )
            ],
            1..1000,
        ),
        // Guaranteed cross-subtree entries
        generate_cross_subtree_entries(),
        // Mix of both (combine random and cross-subtree entries)
        (
            generate_cross_subtree_entries(),
            prop::collection::vec(
                (
                    prop::array::uniform4(arb_felt()).prop_map(RpoDigest::new),
                    prop::array::uniform4(arb_felt())
                ),
                1..1000,
            )
        )
            .prop_map(|(mut cross_subtree, mut random)| {
                cross_subtree.append(&mut random);
                cross_subtree
            })
    ]
    .prop_map(|entries| {
        // Ensure uniqueness of entries as `Smt::with_entries` returns an error if multiple values
        // exist for the same key.
        let mut used_indices = BTreeSet::new();
        let mut used_keys = BTreeSet::new();
        let mut result = Vec::new();

        for (key, value) in entries {
            let leaf_index = LeafIndex::<SMT_DEPTH>::from(key).value();
            if used_indices.insert(leaf_index) && used_keys.insert(key) {
                result.push((key, value));
            }
        }
        result
    })
}

proptest! {
    #[test]
    fn test_with_entries_consistency(entries in arb_entries()) {
        let sequential = Smt::with_entries_sequential(entries.clone()).unwrap();
        let concurrent = Smt::with_entries(entries.clone()).unwrap();
        prop_assert_eq!(concurrent, sequential);
    }

    #[test]
    fn test_compute_mutations_consistency(
        initial_entries in arb_entries(),
        update_entries in arb_entries().prop_filter(
            "Update must not be empty and must differ from initial entries",
            |updates| !updates.is_empty()
        )
    ) {
        let tree = Smt::with_entries_sequential(initial_entries.clone()).unwrap();

        let has_real_changes = update_entries.iter().any(|(key, value)| {
            match initial_entries.iter().find(|(init_key, _)| init_key == key) {
                Some((_, init_value)) => init_value != value,
                None => true,
            }
        });

        let sequential = tree.compute_mutations_sequential(update_entries.clone());
        let concurrent = tree.compute_mutations(update_entries.clone());

        // If there are real changes, the root should change
        if has_real_changes {
            let sequential_changed = sequential.old_root != sequential.new_root;
            let concurrent_changed = concurrent.old_root != concurrent.new_root;

            prop_assert!(
                sequential_changed || concurrent_changed,
                "Root should have changed"
            );
        }

        prop_assert_eq!(sequential.old_root, concurrent.old_root);
        prop_assert_eq!(sequential.new_root, concurrent.new_root);
        prop_assert_eq!(sequential.node_mutations(), concurrent.node_mutations());
        prop_assert_eq!(sequential.new_pairs.len(), concurrent.new_pairs.len());
    }
}