enki 0.0.1

/// Dynamic Fenwick Tree supporting insertions/removals
#[derive(Debug)]
pub struct FenwickTree {
    tree: Vec<usize>, // partial sums
}

impl FenwickTree {
    pub fn new() -> Self {
        Self { tree: vec![0] } // 1-based dummy at index 0
    }

    /// Expand the tree by one element
    pub fn push(&mut self, value: usize) {
        self.tree.push(value);
        let mut i = self.tree.len() - 1;
        while i < self.tree.len() {
            let parent = i + (i & (!i + 1));
            if parent < self.tree.len() {
                self.tree[parent] += value;
            }
            i += 1;
        }
    }

    /// Update index i by delta (i is 0-based)
    pub fn update(&mut self, mut i: usize, delta: isize) {
        i += 1;
        while i < self.tree.len() {
            if delta >= 0 {
                self.tree[i] += delta as usize;
            } else {
                self.tree[i] -= (-delta) as usize;
            }
            i += i & (!i + 1);
        }
    }

    /// Prefix sum 0..i
    pub fn prefix_sum(&self, mut i: usize) -> usize {
        i += 1;
        let mut sum = 0;
        while i > 0 {
            sum += self.tree[i];
            i -= i & (!i + 1);
        }
        sum
    }

    pub fn total_sum(&self) -> usize {
        self.prefix_sum(self.tree.len() - 2)
    }

    /// Find smallest index with cumulative sum >= target
    pub fn lower_bound(&self, target: usize) -> Option<usize> {
        let mut sum = 0;
        let mut pos = 0;
        let n = self.tree.len() - 1;
        let mut k = 1;
        while k < n {
            k <<= 1;
        }

        while k > 0 {
            if pos + k < self.tree.len() && sum + self.tree[pos + k] < target {
                sum += self.tree[pos + k];
                pos += k;
            }
            k >>= 1;
        }

        if pos < n { Some(pos) } else { None }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rand::Rng;

    #[test]
    fn test_push_and_prefix_sum() {
        let mut fw = FenwickTree::new();
        fw.push(3);
        fw.push(1);
        fw.push(4);
        fw.push(2);
        assert_eq!(fw.prefix_sum(0), 3);
        assert_eq!(fw.prefix_sum(1), 4);
        assert_eq!(fw.prefix_sum(2), 8);
        assert_eq!(fw.prefix_sum(3), 10);
    }

    #[test]
    fn test_update_and_total_sum() {
        let mut fw = FenwickTree::new();
        fw.push(2);
        fw.push(5);
        fw.push(7);
        fw.update(1, -2); // 5 -> 3
        fw.update(2, 5); // 7 -> 12
        assert_eq!(fw.prefix_sum(0), 2);
        assert_eq!(fw.prefix_sum(1), 5); // 2 + 3
        assert_eq!(fw.prefix_sum(2), 17); // 2 + 3 + 12
        assert_eq!(fw.total_sum(), 17);
    }

    #[test]
    fn test_lower_bound() {
        let mut fw = FenwickTree::new();
        fw.push(2);
        fw.push(3);
        fw.push(5);
        fw.push(4);
        // prefix sums: [2, 5, 10, 14]
        assert_eq!(fw.lower_bound(1), Some(0));
        assert_eq!(fw.lower_bound(2), Some(0));
        assert_eq!(fw.lower_bound(3), Some(1));
        assert_eq!(fw.lower_bound(6), Some(2));
        assert_eq!(fw.lower_bound(11), Some(3));
        assert_eq!(fw.lower_bound(15), None);
    }

    #[test]
    fn test_prefix_sum_and_total() {
        let mut fw = FenwickTree::new();
        fw.push(2);
        fw.push(3);
        fw.push(5);

        assert_eq!(fw.prefix_sum(0), 2);
        assert_eq!(fw.prefix_sum(1), 5);
        assert_eq!(fw.prefix_sum(2), 10);
        assert_eq!(fw.total_sum(), 10);
    }

    #[test]
    fn test_update_increments_and_decrements() {
        let mut fw = FenwickTree::new();
        fw.push(1);
        fw.push(1);
        fw.push(1);

        assert_eq!(fw.total_sum(), 3);

        fw.update(1, 2); // add 2 at index 1
        assert_eq!(fw.total_sum(), 5);
        assert_eq!(fw.prefix_sum(1), 3);

        fw.update(0, -1); // remove 1 at index 0
        assert_eq!(fw.total_sum(), 4);
        assert_eq!(fw.prefix_sum(0), 0);
    }

    #[test]
    fn test_lower_bound_basic() {
        let mut fw = FenwickTree::new();
        fw.push(2); // indices: 0:2, 1:3, 2:5
        fw.push(3);
        fw.push(5);

        // cumulative: [2, 5, 10]
        assert_eq!(fw.lower_bound(1), Some(0)); // falls in first slot
        assert_eq!(fw.lower_bound(2), Some(0));
        assert_eq!(fw.lower_bound(3), Some(1));
        assert_eq!(fw.lower_bound(5), Some(1));
        assert_eq!(fw.lower_bound(6), Some(2));
        assert_eq!(fw.lower_bound(10), Some(2));
        assert_eq!(fw.lower_bound(11), None); // out of range
    }

    #[test]
    fn test_empty_tree() {
        let fw = FenwickTree::new();
        assert_eq!(fw.total_sum(), 0);
        assert_eq!(fw.lower_bound(1), None);
    }

    #[test]
    fn test_single_element_tree() {
        let mut fw = FenwickTree::new();
        fw.push(7);
        assert_eq!(fw.total_sum(), 7);
        assert_eq!(fw.lower_bound(1), Some(0));
        assert_eq!(fw.lower_bound(7), Some(0));
        assert_eq!(fw.lower_bound(8), None);
    }

    #[test]
    fn test_all_zero_weights() {
        let mut fw = FenwickTree::new();
        fw.push(0);
        fw.push(0);
        fw.push(0);
        assert_eq!(fw.total_sum(), 0);
        assert_eq!(fw.lower_bound(1), None);
    }

    #[test]
    fn test_weighted_sampling_distribution() {
        let mut rng = rand::rng();
        let mut fw = FenwickTree::new();
        fw.push(1);
        fw.push(3);
        fw.push(6);

        let mut counts = vec![0, 0, 0];
        let trials = 10_000;

        for _ in 0..trials {
            let total = fw.total_sum();
            let r = rng.random_range(1..=total);
            let idx = fw.lower_bound(r).unwrap();
            counts[idx] += 1;
        }

        // Expect approx 1:3:6 ratio
        let ratio: Vec<f64> = counts.iter().map(|&c| c as f64 / trials as f64).collect();
        assert!((ratio[0] - 0.1).abs() < 0.05);
        assert!((ratio[1] - 0.3).abs() < 0.05);
        assert!((ratio[2] - 0.6).abs() < 0.05);
    }
}