numrs2 0.3.3

//! Tests for newly implemented NumPy parity functions
//!
//! This module tests the new functions implemented for NumPy parity:
//! - take() and take_along_axis() for advanced indexing
//! - Set operations: intersect1d, union1d, setdiff1d, setxor1d, in1d, isin
//! - Conditional operations: where_cond, select

use approx::assert_relative_eq;
use numrs2::math::{linspace, meshgrid, mgrid, ogrid};
use numrs2::prelude::*;

#[cfg(test)]
mod test_indexing_functions {
    use super::*;

    #[test]
    fn test_take_flattened() {
        let a = Array::from_vec(vec![10, 20, 30, 40, 50]);
        let indices = Array::from_vec(vec![0usize, 2, 4]);
        let result = take(&a, &indices, None, None).unwrap();
        assert_eq!(result.to_vec(), vec![10, 30, 50]);
    }

    #[test]
    fn test_take_along_axis() {
        let a = Array::from_vec(vec![1, 2, 3, 4, 5, 6]).reshape(&[2, 3]);
        let indices = Array::from_vec(vec![0usize, 2]);
        let result = take(&a, &indices, Some(1), None).unwrap();
        assert_eq!(result.shape(), vec![2, 2]);
        // The result order depends on how the function iterates through the data
        // Let's check the actual result and adjust
        let actual = result.to_vec();
        // Should contain elements [1, 3] (first row indices 0,2) and [4, 6] (second row indices 0,2)
        assert!(
            actual.contains(&1)
                && actual.contains(&3)
                && actual.contains(&4)
                && actual.contains(&6)
        );
        assert_eq!(actual.len(), 4);
    }

    #[test]
    fn test_take_wrap_mode() {
        let a = Array::from_vec(vec![10, 20, 30]);
        let indices = Array::from_vec(vec![0usize, 1, 2, 3, 4, 5]);
        let result = take(&a, &indices, None, Some("wrap")).unwrap();
        assert_eq!(result.to_vec(), vec![10, 20, 30, 10, 20, 30]);
    }

    #[test]
    fn test_take_clip_mode() {
        let a = Array::from_vec(vec![10, 20, 30]);
        let indices = Array::from_vec(vec![0usize, 1, 2, 3, 4]);
        let result = take(&a, &indices, None, Some("clip")).unwrap();
        assert_eq!(result.to_vec(), vec![10, 20, 30, 30, 30]);
    }

    #[test]
    fn test_take_along_axis_2d() {
        let a = Array::from_vec(vec![10, 20, 30, 40, 50, 60]).reshape(&[2, 3]);
        let indices = Array::from_vec(vec![2usize, 0, 1, 1]).reshape(&[2, 2]);
        let result = take_along_axis(&a, &indices, 1).unwrap();
        assert_eq!(result.shape(), vec![2, 2]);
        assert_eq!(result.to_vec(), vec![30, 10, 50, 50]);
    }

    #[test]
    fn test_take_error_invalid_index() {
        let a = Array::from_vec(vec![10, 20, 30]);
        let indices = Array::from_vec(vec![0usize, 1, 5]);
        let result = take(&a, &indices, None, None);
        assert!(result.is_err());
    }
}

#[cfg(test)]
mod test_set_operations {
    use super::*;

    #[test]
    fn test_intersect1d() {
        let a = Array::from_vec(vec![1, 3, 4, 3]);
        let b = Array::from_vec(vec![3, 1, 2, 1]);
        let result = intersect1d(&a, &b).unwrap();
        assert_eq!(result.to_vec(), vec![1, 3]);
    }

    #[test]
    fn test_union1d() {
        let a = Array::from_vec(vec![1, 2, 3, 2, 4]);
        let b = Array::from_vec(vec![2, 3, 5, 7, 5]);
        let result = union1d(&a, &b).unwrap();
        assert_eq!(result.to_vec(), vec![1, 2, 3, 4, 5, 7]);
    }

    #[test]
    fn test_setdiff1d() {
        let a = Array::from_vec(vec![1, 2, 3, 2, 4, 1]);
        let b = Array::from_vec(vec![3, 4, 5, 6]);
        let result = setdiff1d(&a, &b).unwrap();
        assert_eq!(result.to_vec(), vec![1, 2]);
    }

    #[test]
    fn test_setxor1d() {
        let a = Array::from_vec(vec![1, 2, 3, 2, 4]);
        let b = Array::from_vec(vec![2, 3, 5, 7, 5]);
        let result = setxor1d(&a, &b).unwrap();
        assert_eq!(result.to_vec(), vec![1, 4, 5, 7]);
    }

    #[test]
    fn test_in1d() {
        let a = Array::from_vec(vec![1, 2, 3, 4, 5, 6]);
        let b = Array::from_vec(vec![2, 4, 6]);
        let result = in1d(&a, &b).unwrap();
        assert_eq!(result.to_vec(), vec![false, true, false, true, false, true]);
    }

    #[test]
    fn test_isin() {
        let element = Array::from_vec(vec![0, 1, 2, 5, 0]);
        let test_elements = Array::from_vec(vec![0, 2, 5, 7, 9]);
        let result = isin(&element, &test_elements, false, false).unwrap();
        assert_eq!(result.to_vec(), vec![true, false, true, true, true]);

        // With invert=true
        let result_inv = isin(&element, &test_elements, false, true).unwrap();
        assert_eq!(result_inv.to_vec(), vec![false, true, false, false, false]);
    }

    #[test]
    fn test_isin_2d() {
        let element = Array::from_vec(vec![0, 1, 2, 3]).reshape(&[2, 2]);
        let test_elements = Array::from_vec(vec![0, 2]);
        let result = isin(&element, &test_elements, false, false).unwrap();
        assert_eq!(result.shape(), vec![2, 2]);
        assert_eq!(result.to_vec(), vec![true, false, true, false]);
    }

    #[test]
    fn test_unique_with_options() {
        let a = Array::from_vec(vec![1, 1, 2, 2, 3, 3]);
        let (unique, indices, inverse, counts) = unique_with_options(&a, true, true, true).unwrap();
        assert_eq!(unique.to_vec(), vec![1, 2, 3]);
        assert_eq!(indices.to_vec(), vec![0, 2, 4]);
        assert_eq!(inverse.to_vec(), vec![0, 0, 1, 1, 2, 2]);
        assert_eq!(counts.to_vec(), vec![2, 2, 2]);
    }

    #[test]
    fn test_unique_with_options_complex() {
        let a = Array::from_vec(vec![3, 1, 2, 1, 3, 2, 1]);
        let (unique, indices, inverse, counts) = unique_with_options(&a, true, true, true).unwrap();
        assert_eq!(unique.to_vec(), vec![1, 2, 3]);
        assert_eq!(indices.to_vec(), vec![1, 2, 0]); // First occurrence indices
        assert_eq!(inverse.to_vec(), vec![2, 0, 1, 0, 2, 1, 0]); // Reconstruct indices
        assert_eq!(counts.to_vec(), vec![3, 2, 2]); // Count of each unique value
    }

    #[test]
    fn test_empty_arrays() {
        let a: Array<i32> = Array::from_vec(vec![]);
        let b: Array<i32> = Array::from_vec(vec![]);

        let intersection = intersect1d(&a, &b).unwrap();
        assert_eq!(intersection.to_vec(), Vec::<i32>::new());

        let union = union1d(&a, &b).unwrap();
        assert_eq!(union.to_vec(), Vec::<i32>::new());
    }
}

#[cfg(test)]
mod test_conditional_operations {
    use super::*;

    #[test]
    fn test_where_cond_basic() {
        let condition = Array::from_vec(vec![true, false, true, false]);
        let x = Array::from_vec(vec![1, 2, 3, 4]);
        let y = Array::from_vec(vec![10, 20, 30, 40]);
        let result = where_cond(&condition, &x, &y).unwrap();
        assert_eq!(result.to_vec(), vec![1, 20, 3, 40]);
    }

    #[test]
    fn test_where_cond_broadcasting() {
        let condition = Array::from_vec(vec![true, false, true, false]).reshape(&[2, 2]);
        let x = Array::from_vec(vec![100]);
        let y = Array::from_vec(vec![1, 2, 3, 4]).reshape(&[2, 2]);
        let result = where_cond(&condition, &x, &y).unwrap();
        assert_eq!(result.to_vec(), vec![100, 2, 100, 4]);
        assert_eq!(result.shape(), vec![2, 2]);
    }

    #[test]
    fn test_where_cond_same_arrays() {
        let condition = Array::from_vec(vec![false, false, false]);
        let x = Array::from_vec(vec![1, 2, 3]);
        let y = Array::from_vec(vec![10, 20, 30]);
        let result = where_cond(&condition, &x, &y).unwrap();
        assert_eq!(result.to_vec(), vec![10, 20, 30]);
    }

    #[test]
    fn test_select_basic() {
        // Create conditions
        let x = Array::from_vec(vec![0, 1, 2, 3, 4, 5]);
        let cond1 = x.map(|val| val < 3);
        let cond2 = x.map(|val| val >= 3);

        // Create choices
        let choice1 = Array::from_vec(vec![10, 10, 10, 10, 10, 10]);
        let choice2 = Array::from_vec(vec![20, 20, 20, 20, 20, 20]);

        let result = select(&[&cond1, &cond2], &[&choice1, &choice2], Some(99)).unwrap();
        assert_eq!(result.to_vec(), vec![10, 10, 10, 20, 20, 20]);
    }

    #[test]
    fn test_select_default() {
        // When no condition matches, use default
        let always_false = Array::from_vec(vec![false, false, false]);
        let choice_unused = Array::from_vec(vec![1, 2, 3]);
        let result = select(&[&always_false], &[&choice_unused], Some(99)).unwrap();
        assert_eq!(result.to_vec(), vec![99, 99, 99]);
    }

    #[test]
    fn test_select_overlapping_conditions() {
        // Test with overlapping conditions - first match wins
        let always_true = Array::from_vec(vec![true, true, true]);
        let also_always_true = Array::from_vec(vec![true, true, true]);
        let choice1 = Array::from_vec(vec![10, 10, 10]);
        let choice2 = Array::from_vec(vec![20, 20, 20]);

        let result = select(
            &[&always_true, &also_always_true],
            &[&choice1, &choice2],
            Some(99),
        )
        .unwrap();
        assert_eq!(result.to_vec(), vec![10, 10, 10]); // First condition wins
    }

    #[test]
    fn test_select_error_mismatched_lengths() {
        let cond1 = Array::from_vec(vec![true, false]);
        let choice1 = Array::from_vec(vec![1, 2]);
        let choice2 = Array::from_vec(vec![3, 4]);

        // Mismatch: 1 condition, 2 choices
        let result = select(&[&cond1], &[&choice1, &choice2], Some(0));
        assert!(result.is_err());
    }
}

#[cfg(test)]
mod test_grid_functions {
    use super::*;

    #[test]
    fn test_meshgrid_basic() {
        let x = Array::from_vec(vec![1.0, 2.0, 3.0]);
        let y = Array::from_vec(vec![4.0, 5.0, 6.0, 7.0]);

        // xy indexing (default)
        let grids = meshgrid(&[&x, &y], None).unwrap();
        let xx = &grids[0];
        let yy = &grids[1];
        assert_eq!(xx.shape(), vec![4, 3]); // (y.len(), x.len())
        assert_eq!(yy.shape(), vec![4, 3]);

        // Check specific values for xy indexing
        assert_eq!(xx.array().get([0, 0]).unwrap(), &1.0);
        assert_eq!(xx.array().get([0, 1]).unwrap(), &2.0);
        assert_eq!(xx.array().get([0, 2]).unwrap(), &3.0);
        assert_eq!(yy.array().get([0, 0]).unwrap(), &4.0);
        assert_eq!(yy.array().get([1, 0]).unwrap(), &5.0);
    }

    #[test]
    fn test_meshgrid_ij_indexing() {
        let x = Array::from_vec(vec![1.0, 2.0, 3.0]);
        let y = Array::from_vec(vec![4.0, 5.0, 6.0, 7.0]);

        // ij indexing
        let grids = meshgrid(&[&x, &y], Some("ij")).unwrap();
        let ii = &grids[0];
        let jj = &grids[1];
        assert_eq!(ii.shape(), vec![3, 4]); // (x.len(), y.len())
        assert_eq!(jj.shape(), vec![3, 4]);

        // Check specific values for ij indexing
        assert_eq!(ii.array().get([0, 0]).unwrap(), &1.0);
        assert_eq!(ii.array().get([1, 0]).unwrap(), &2.0);
        assert_eq!(ii.array().get([2, 0]).unwrap(), &3.0);
        assert_eq!(jj.array().get([0, 0]).unwrap(), &4.0);
        assert_eq!(jj.array().get([0, 1]).unwrap(), &5.0);
    }

    #[test]
    fn test_mgrid_basic() {
        let x = linspace(0.0, 2.0, 3);
        let y = linspace(0.0, 2.0, 3);

        let grids = mgrid(&[&x, &y]).unwrap();
        let xx = &grids[0];
        let yy = &grids[1];

        assert_eq!(xx.shape(), vec![3, 3]);
        assert_eq!(yy.shape(), vec![3, 3]);

        // Check some values
        assert_eq!(xx.array().get([0, 0]).unwrap(), &0.0);
        assert_eq!(xx.array().get([0, 1]).unwrap(), &0.0);
        assert_eq!(xx.array().get([1, 0]).unwrap(), &1.0);
        assert_eq!(yy.array().get([0, 0]).unwrap(), &0.0);
        assert_eq!(yy.array().get([1, 0]).unwrap(), &0.0);
        assert_eq!(yy.array().get([0, 1]).unwrap(), &1.0);
    }

    #[test]
    fn test_ogrid_basic() {
        let x = linspace(0.0, 2.0, 3);
        let y = linspace(0.0, 2.0, 3);

        let grids = ogrid(&[&x, &y]).unwrap();
        let xx = &grids[0];
        let yy = &grids[1];

        assert_eq!(xx.shape(), vec![3, 1]);
        assert_eq!(yy.shape(), vec![1, 3]);

        // Check values
        assert_eq!(xx.array().get([0, 0]).unwrap(), &0.0);
        assert_eq!(xx.array().get([1, 0]).unwrap(), &1.0);
        assert_eq!(xx.array().get([2, 0]).unwrap(), &2.0);

        assert_eq!(yy.array().get([0, 0]).unwrap(), &0.0);
        assert_eq!(yy.array().get([0, 1]).unwrap(), &1.0);
        assert_eq!(yy.array().get([0, 2]).unwrap(), &2.0);
    }
}

#[cfg(test)]
mod test_edge_cases {
    use super::*;

    #[test]
    fn test_empty_inputs() {
        let empty: Array<i32> = Array::from_vec(vec![]);
        let indices: Array<usize> = Array::from_vec(vec![]);
        let result = take(&empty, &indices, None, None).unwrap();
        assert_eq!(result.to_vec(), Vec::<i32>::new());
    }

    #[test]
    fn test_single_element_arrays() {
        let a = Array::from_vec(vec![42]);
        let indices = Array::from_vec(vec![0usize]);
        let result = take(&a, &indices, None, None).unwrap();
        assert_eq!(result.to_vec(), vec![42]);
    }

    #[test]
    fn test_where_cond_all_true() {
        let condition = Array::from_vec(vec![true, true, true]);
        let x = Array::from_vec(vec![1, 2, 3]);
        let y = Array::from_vec(vec![10, 20, 30]);
        let result = where_cond(&condition, &x, &y).unwrap();
        assert_eq!(result.to_vec(), vec![1, 2, 3]);
    }

    #[test]
    fn test_where_cond_all_false() {
        let condition = Array::from_vec(vec![false, false, false]);
        let x = Array::from_vec(vec![1, 2, 3]);
        let y = Array::from_vec(vec![10, 20, 30]);
        let result = where_cond(&condition, &x, &y).unwrap();
        assert_eq!(result.to_vec(), vec![10, 20, 30]);
    }
}

#[cfg(test)]
mod test_new_indexing_functions {
    use super::*;

    #[test]
    fn test_put_along_axis_basic() {
        let mut a = Array::zeros(&[2, 3]);
        let indices = Array::from_vec(vec![2, 0, 1, 1]).reshape(&[2, 2]);
        let values = Array::from_vec(vec![10, 20, 30, 40]).reshape(&[2, 2]);
        put_along_axis(&mut a, &indices, &values, 1).unwrap();

        assert_eq!(a.get(&[0, 2]).unwrap(), 10);
        assert_eq!(a.get(&[0, 0]).unwrap(), 20);
        assert_eq!(a.get(&[1, 1]).unwrap(), 40); // Last value overwrites
    }

    #[test]
    fn test_put_along_axis_1d() {
        let mut a = Array::zeros(&[5]);
        let indices = Array::from_vec(vec![0, 2, 4]);
        let values = Array::from_vec(vec![10, 20, 30]);
        put_along_axis(&mut a, &indices, &values, 0).unwrap();

        assert_eq!(a.to_vec(), vec![10, 0, 20, 0, 30]);
    }

    #[test]
    fn test_put_along_axis_error_bounds() {
        let mut a = Array::zeros(&[2, 3]);
        let indices = Array::from_vec(vec![5]); // Out of bounds
        let values = Array::from_vec(vec![10]);
        let result = put_along_axis(&mut a, &indices, &values, 1);
        assert!(result.is_err());
    }

    #[test]
    fn test_extract_basic() {
        let a = Array::from_vec(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
        let condition = a.map(|x| x > 5);
        let result = extract(&a, &condition).unwrap();
        assert_eq!(result.to_vec(), vec![6, 7, 8, 9, 10]);
    }

    #[test]
    fn test_extract_2d() {
        let a = Array::from_vec(vec![1, 2, 3, 4, 5, 6]).reshape(&[2, 3]);
        let condition = a.map(|x| x % 2 == 0);
        let result = extract(&a, &condition).unwrap();
        assert_eq!(result.to_vec(), vec![2, 4, 6]);
    }

    #[test]
    fn test_extract_empty() {
        let a = Array::from_vec(vec![1, 2, 3, 4, 5]);
        let condition = a.map(|x| x > 10); // No elements satisfy
        let result = extract(&a, &condition).unwrap();
        assert_eq!(result.to_vec(), Vec::<i32>::new());
    }

    #[test]
    fn test_extract_all() {
        let a = Array::from_vec(vec![1, 2, 3, 4, 5]);
        let condition = a.map(|_x| true); // All elements satisfy
        let result = extract(&a, &condition).unwrap();
        assert_eq!(result.to_vec(), vec![1, 2, 3, 4, 5]);
    }

    #[test]
    fn test_extract_shape_mismatch() {
        let a = Array::from_vec(vec![1, 2, 3, 4]);
        let condition = Array::from_vec(vec![true, false]); // Wrong shape
        let result = extract(&a, &condition);
        assert!(result.is_err());
    }
}

#[cfg(test)]
mod test_fromfunction {
    use super::*;

    #[test]
    fn test_fromfunction_basic_2d() {
        let result = fromfunction(
            |indices: &[usize]| (indices[0] + indices[1]) as f64,
            &[3, 3],
        )
        .unwrap();
        assert_eq!(result.get(&[0, 0]).unwrap(), 0.0);
        assert_eq!(result.get(&[0, 1]).unwrap(), 1.0);
        assert_eq!(result.get(&[1, 0]).unwrap(), 1.0);
        assert_eq!(result.get(&[1, 1]).unwrap(), 2.0);
        assert_eq!(result.get(&[2, 2]).unwrap(), 4.0);
    }

    #[test]
    fn test_fromfunction_multiplication() {
        let result = fromfunction(
            |indices: &[usize]| (indices[0] * indices[1]) as i32,
            &[3, 4],
        )
        .unwrap();
        assert_eq!(result.get(&[0, 0]).unwrap(), 0);
        assert_eq!(result.get(&[0, 3]).unwrap(), 0);
        assert_eq!(result.get(&[1, 3]).unwrap(), 3);
        assert_eq!(result.get(&[2, 3]).unwrap(), 6);
    }

    #[test]
    fn test_fromfunction_1d() {
        let result = fromfunction(|indices: &[usize]| indices[0] as f64, &[5]).unwrap();
        assert_eq!(result.to_vec(), vec![0.0, 1.0, 2.0, 3.0, 4.0]);
    }

    #[test]
    fn test_fromfunction_3d() {
        let result = fromfunction(
            |indices: &[usize]| (indices[0] + indices[1] + indices[2]) as i32,
            &[2, 2, 2],
        )
        .unwrap();
        assert_eq!(result.get(&[0, 0, 0]).unwrap(), 0);
        assert_eq!(result.get(&[1, 1, 1]).unwrap(), 3);
        assert_eq!(result.shape(), vec![2, 2, 2]);
    }

    #[test]
    fn test_fromfunction_distance() {
        // Create a distance-from-center function
        let result = fromfunction(
            |indices: &[usize]| {
                let x = indices[0] as f64 - 1.0; // Center at (1, 1)
                let y = indices[1] as f64 - 1.0;
                (x * x + y * y).sqrt()
            },
            &[3, 3],
        )
        .unwrap();

        assert_eq!(result.get(&[1, 1]).unwrap(), 0.0); // Center
        assert_relative_eq!(result.get(&[0, 1]).unwrap(), 1.0, epsilon = 1e-10);
        assert_relative_eq!(
            result.get(&[0, 0]).unwrap(),
            std::f64::consts::SQRT_2,
            epsilon = 1e-10
        );
    }

    #[test]
    fn test_fromfunction_empty() {
        let result = fromfunction(|_indices: &[usize]| 42i32, &[]).unwrap();
        assert_eq!(result.to_vec(), Vec::<i32>::new());
    }
}

#[cfg(test)]
mod test_unique_axis {
    use super::*;

    #[test]
    fn test_unique_axis_none() {
        // Test with axis=None (should behave like unique_with_options)
        let a = Array::from_vec(vec![3, 1, 2, 1, 3, 2, 1]);
        let (unique, indices, inverse, counts) = unique_axis(&a, None, true, true, true).unwrap();
        assert_eq!(unique.to_vec(), vec![1, 2, 3]);
        assert_eq!(indices.to_vec(), vec![1, 2, 0]);
        assert_eq!(inverse.to_vec(), vec![2, 0, 1, 0, 2, 1, 0]);
        assert_eq!(counts.to_vec(), vec![3, 2, 2]);
    }

    #[test]
    fn test_unique_axis_rows() {
        // Test finding unique rows (axis=0)
        let a = Array::from_vec(vec![1, 2, 3, 1, 2, 3, 4, 5, 6]).reshape(&[3, 3]);
        let (unique_rows, indices, inverse, counts) =
            unique_axis(&a, Some(0), true, true, true).unwrap();

        assert_eq!(unique_rows.shape(), vec![2, 3]); // Two unique rows
        assert_eq!(indices.to_vec(), vec![0, 2]); // First occurrences
        assert_eq!(inverse.to_vec(), vec![0, 0, 1]); // Row 0 and 1 are the same, row 2 is different
        assert_eq!(counts.to_vec(), vec![2, 1]); // First row appears twice, second row appears once
    }

    #[test]
    fn test_unique_axis_columns() {
        // Test finding unique columns (axis=1)
        let a = Array::from_vec(vec![1, 2, 1, 3, 4, 3]).reshape(&[2, 3]);
        let (unique_cols, indices, inverse, counts) =
            unique_axis(&a, Some(1), true, true, true).unwrap();

        assert_eq!(unique_cols.shape(), vec![2, 2]); // Two unique columns
        assert_eq!(indices.to_vec(), vec![0, 1]); // First occurrences
        assert_eq!(inverse.to_vec(), vec![0, 1, 0]); // Col 0 and 2 are the same, col 1 is different
        assert_eq!(counts.to_vec(), vec![2, 1]); // First column appears twice, second column appears once
    }

    #[test]
    fn test_unique_axis_all_unique() {
        // Test when all rows are unique
        let a = Array::from_vec(vec![1, 2, 3, 4, 5, 6]).reshape(&[2, 3]);
        let (unique_rows, indices, inverse, counts) =
            unique_axis(&a, Some(0), true, true, true).unwrap();

        assert_eq!(unique_rows.shape(), vec![2, 3]); // All rows are unique
        assert_eq!(indices.to_vec(), vec![0, 1]);
        assert_eq!(inverse.to_vec(), vec![0, 1]);
        assert_eq!(counts.to_vec(), vec![1, 1]);
    }

    #[test]
    fn test_unique_axis_single_row() {
        // Test with single row
        let a = Array::from_vec(vec![1, 2, 3]).reshape(&[1, 3]);
        let (unique_rows, indices, inverse, counts) =
            unique_axis(&a, Some(0), true, true, true).unwrap();

        assert_eq!(unique_rows.shape(), vec![1, 3]);
        assert_eq!(indices.to_vec(), vec![0]);
        assert_eq!(inverse.to_vec(), vec![0]);
        assert_eq!(counts.to_vec(), vec![1]);
    }

    #[test]
    fn test_unique_axis_error_invalid_axis() {
        let a = Array::from_vec(vec![1, 2, 3, 4]).reshape(&[2, 2]);
        let result = unique_axis(&a, Some(2), false, false, false); // axis=2 is invalid for 2D array
        assert!(result.is_err());
    }

    #[test]
    fn test_unique_axis_no_return_options() {
        // Test with all return options disabled
        let a = Array::from_vec(vec![1, 2, 1, 3, 4, 3]).reshape(&[2, 3]);
        let (unique_cols, indices, inverse, counts) =
            unique_axis(&a, Some(1), false, false, false).unwrap();

        assert_eq!(unique_cols.shape(), vec![2, 2]); // Two unique columns
        assert_eq!(indices.to_vec(), Vec::<usize>::new()); // Empty
        assert_eq!(inverse.to_vec(), Vec::<usize>::new()); // Empty
        assert_eq!(counts.to_vec(), Vec::<usize>::new()); // Empty
    }
}

#[cfg(test)]
mod test_enhanced_eye_identity {
    use super::*;

    #[test]
    fn test_enhanced_eye_basic() {
        let eye = Array::<i32>::eye(3, 3, 0);
        assert_eq!(eye.shape(), vec![3, 3]);
        assert_eq!(eye.to_vec(), vec![1, 0, 0, 0, 1, 0, 0, 0, 1]);
    }

    #[test]
    fn test_enhanced_eye_rectangular() {
        let eye = Array::<f64>::eye(2, 4, 0);
        assert_eq!(eye.shape(), vec![2, 4]);
        assert_eq!(eye.to_vec(), vec![1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0]);
    }

    #[test]
    fn test_enhanced_eye_diagonal_above() {
        let eye = Array::<i32>::eye(3, 3, 1);
        assert_eq!(eye.shape(), vec![3, 3]);
        assert_eq!(eye.to_vec(), vec![0, 1, 0, 0, 0, 1, 0, 0, 0]);
    }

    #[test]
    fn test_enhanced_eye_diagonal_below() {
        let eye = Array::<i32>::eye(3, 3, -1);
        assert_eq!(eye.shape(), vec![3, 3]);
        assert_eq!(eye.to_vec(), vec![0, 0, 0, 1, 0, 0, 0, 1, 0]);
    }

    #[test]
    fn test_enhanced_eye_large_offset() {
        // Test with diagonal that goes beyond matrix bounds
        let eye = Array::<i32>::eye(3, 3, 5);
        assert_eq!(eye.shape(), vec![3, 3]);
        assert_eq!(eye.to_vec(), vec![0, 0, 0, 0, 0, 0, 0, 0, 0]); // All zeros
    }

    #[test]
    fn test_enhanced_eye_negative_large_offset() {
        // Test with negative diagonal that goes beyond matrix bounds
        let eye = Array::<i32>::eye(3, 3, -5);
        assert_eq!(eye.shape(), vec![3, 3]);
        assert_eq!(eye.to_vec(), vec![0, 0, 0, 0, 0, 0, 0, 0, 0]); // All zeros
    }

    #[test]
    fn test_identity_basic() {
        let identity = Array::<f64>::identity(4);
        assert_eq!(identity.shape(), vec![4, 4]);
        assert_eq!(
            identity.to_vec(),
            vec![1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0]
        );
    }
}

#[cfg(test)]
mod test_frombuffer_fromiter {
    use super::*;

    #[test]
    fn test_frombuffer_i32() {
        let data: Vec<i32> = vec![1, 2, 3, 4, 5];
        let buffer = unsafe {
            std::slice::from_raw_parts(
                data.as_ptr() as *const u8,
                data.len() * std::mem::size_of::<i32>(),
            )
        };
        let result = frombuffer::<i32>(buffer, std::mem::size_of::<i32>(), -1, 0).unwrap();
        assert_eq!(result.to_vec(), vec![1, 2, 3, 4, 5]);
    }

    #[test]
    fn test_frombuffer_with_count() {
        let data: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let buffer = unsafe {
            std::slice::from_raw_parts(
                data.as_ptr() as *const u8,
                data.len() * std::mem::size_of::<f64>(),
            )
        };
        let result = frombuffer::<f64>(buffer, std::mem::size_of::<f64>(), 3, 0).unwrap();
        assert_eq!(result.to_vec(), vec![1.0, 2.0, 3.0]);
    }

    #[test]
    fn test_frombuffer_with_offset() {
        let data: Vec<i32> = vec![1, 2, 3, 4, 5];
        let buffer = unsafe {
            std::slice::from_raw_parts(
                data.as_ptr() as *const u8,
                data.len() * std::mem::size_of::<i32>(),
            )
        };
        let offset = 2 * std::mem::size_of::<i32>(); // Skip first 2 elements
        let result = frombuffer::<i32>(buffer, std::mem::size_of::<i32>(), -1, offset).unwrap();
        assert_eq!(result.to_vec(), vec![3, 4, 5]);
    }

    #[test]
    fn test_frombuffer_error_invalid_size() {
        let data: Vec<i32> = vec![1, 2, 3];
        let buffer = unsafe {
            std::slice::from_raw_parts(
                data.as_ptr() as *const u8,
                data.len() * std::mem::size_of::<i32>(),
            )
        };
        let result = frombuffer::<i32>(buffer, 8, -1, 0); // Wrong size
        assert!(result.is_err());
    }

    #[test]
    fn test_frombuffer_error_offset_too_large() {
        let data: Vec<i32> = vec![1, 2, 3];
        let buffer = unsafe {
            std::slice::from_raw_parts(
                data.as_ptr() as *const u8,
                data.len() * std::mem::size_of::<i32>(),
            )
        };
        let result = frombuffer::<i32>(buffer, std::mem::size_of::<i32>(), -1, 100);
        assert!(result.is_err());
    }

    #[test]
    fn test_fromiter_1d() {
        let result = fromiter((0..5).map(|x| x as f64), None).unwrap();
        assert_eq!(result.to_vec(), vec![0.0, 1.0, 2.0, 3.0, 4.0]);
    }

    #[test]
    fn test_fromiter_2d() {
        let result = fromiter(0..6, Some(&[2, 3])).unwrap();
        assert_eq!(result.shape(), vec![2, 3]);
        assert_eq!(result.to_vec(), vec![0, 1, 2, 3, 4, 5]);
    }

    #[test]
    fn test_fromiter_shape_mismatch() {
        let result = fromiter(0..5, Some(&[2, 3])); // 5 elements but 2x3=6 expected
        assert!(result.is_err());
    }

    #[test]
    fn test_fromiter_empty() {
        let result = fromiter(std::iter::empty::<i32>(), None).unwrap();
        assert_eq!(result.to_vec(), Vec::<i32>::new());
    }

    #[test]
    fn test_fromiter_complex() {
        // Test with more complex iterator
        let result = fromiter((0..4).map(|x| x * x), Some(&[2, 2])).unwrap();
        assert_eq!(result.shape(), vec![2, 2]);
        assert_eq!(result.to_vec(), vec![0, 1, 4, 9]);
    }
}

#[cfg(test)]
mod test_nan_arg_functions {
    use super::*;
    use numrs2::math::{nanargmax, nanargmin, nancumprod};

    #[test]
    fn test_nanargmax_1d_no_nan() {
        let a = Array::from_vec(vec![1.0, 5.0, 3.0, 2.0]);
        let result = nanargmax(&a, None, false).unwrap();
        assert_eq!(result.to_vec(), vec![1]); // index 1 has value 5.0
    }

    #[test]
    fn test_nanargmax_1d_with_nan() {
        let a = Array::from_vec(vec![1.0, f64::NAN, 3.0, 2.0]);
        let result = nanargmax(&a, None, false).unwrap();
        assert_eq!(result.to_vec(), vec![2]); // index 2 has value 3.0, ignoring NaN
    }

    #[test]
    fn test_nanargmax_1d_nan_at_max() {
        // NaN at what would be max position
        let a = Array::from_vec(vec![1.0, 2.0, f64::NAN, 4.0]);
        let result = nanargmax(&a, None, false).unwrap();
        assert_eq!(result.to_vec(), vec![3]); // index 3 has value 4.0
    }

    #[test]
    fn test_nanargmax_2d_axis0() {
        let a = Array::from_vec(vec![1.0, f64::NAN, 3.0, 4.0, 2.0, 1.0]).reshape(&[2, 3]);
        let result = nanargmax(&a, Some(0), false).unwrap();
        assert_eq!(result.shape(), vec![3]);
        assert_eq!(result.to_vec(), vec![1, 1, 0]); // max indices along axis 0
    }

    #[test]
    fn test_nanargmax_2d_axis1() {
        let a = Array::from_vec(vec![1.0, f64::NAN, 3.0, 2.0, 4.0, 1.0]).reshape(&[2, 3]);
        let result = nanargmax(&a, Some(1), false).unwrap();
        assert_eq!(result.shape(), vec![2]);
        assert_eq!(result.to_vec(), vec![2, 1]); // max indices in each row, ignoring NaN
    }

    #[test]
    fn test_nanargmax_keepdims() {
        let a = Array::from_vec(vec![1.0, f64::NAN, 3.0, 2.0]);
        let result = nanargmax(&a, None, true).unwrap();
        assert_eq!(result.shape(), vec![1]);
        assert_eq!(result.to_vec(), vec![2]);
    }

    #[test]
    fn test_nanargmin_1d_no_nan() {
        let a = Array::from_vec(vec![5.0, 1.0, 3.0, 2.0]);
        let result = nanargmin(&a, None, false).unwrap();
        assert_eq!(result.to_vec(), vec![1]); // index 1 has value 1.0
    }

    #[test]
    fn test_nanargmin_1d_with_nan() {
        let a = Array::from_vec(vec![3.0, f64::NAN, 1.0, 2.0]);
        let result = nanargmin(&a, None, false).unwrap();
        assert_eq!(result.to_vec(), vec![2]); // index 2 has value 1.0, ignoring NaN
    }

    #[test]
    fn test_nanargmin_1d_nan_at_min() {
        // NaN at what would be min position
        let a = Array::from_vec(vec![4.0, 2.0, f64::NAN, 1.0]);
        let result = nanargmin(&a, None, false).unwrap();
        assert_eq!(result.to_vec(), vec![3]); // index 3 has value 1.0
    }

    #[test]
    fn test_nanargmin_2d_axis0() {
        let a = Array::from_vec(vec![4.0, f64::NAN, 1.0, 2.0, 5.0, 3.0]).reshape(&[2, 3]);
        let result = nanargmin(&a, Some(0), false).unwrap();
        assert_eq!(result.shape(), vec![3]);
        assert_eq!(result.to_vec(), vec![1, 1, 0]); // min indices along axis 0
    }

    #[test]
    fn test_nanargmin_2d_axis1() {
        let a = Array::from_vec(vec![3.0, f64::NAN, 1.0, 2.0, 4.0, 0.5]).reshape(&[2, 3]);
        let result = nanargmin(&a, Some(1), false).unwrap();
        assert_eq!(result.shape(), vec![2]);
        assert_eq!(result.to_vec(), vec![2, 2]); // min indices in each row, ignoring NaN
    }

    #[test]
    fn test_nanargmin_keepdims() {
        let a = Array::from_vec(vec![3.0, f64::NAN, 1.0, 2.0]);
        let result = nanargmin(&a, None, true).unwrap();
        assert_eq!(result.shape(), vec![1]);
        assert_eq!(result.to_vec(), vec![2]);
    }

    #[test]
    fn test_nancumprod_1d_no_nan() {
        let a = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0]);
        let result = nancumprod(&a, None).unwrap();
        assert_eq!(result.to_vec(), vec![1.0, 2.0, 6.0, 24.0]);
    }

    #[test]
    fn test_nancumprod_1d_with_nan() {
        let a = Array::from_vec(vec![1.0, f64::NAN, 2.0, 3.0]);
        let result = nancumprod(&a, None).unwrap();
        assert_eq!(result.to_vec(), vec![1.0, 1.0, 2.0, 6.0]); // NaN treated as 1
    }

    #[test]
    fn test_nancumprod_1d_multiple_nan() {
        let a = Array::from_vec(vec![2.0, f64::NAN, f64::NAN, 3.0]);
        let result = nancumprod(&a, None).unwrap();
        assert_eq!(result.to_vec(), vec![2.0, 2.0, 2.0, 6.0]); // Multiple NaNs treated as 1
    }

    #[test]
    fn test_nancumprod_2d_axis0() {
        let a = Array::from_vec(vec![1.0, 2.0, f64::NAN, 3.0, 4.0, 2.0]).reshape(&[2, 3]);
        let result = nancumprod(&a, Some(0)).unwrap();
        assert_eq!(result.shape(), vec![2, 3]);
        // First row: [1.0, 2.0, NaN->1.0]
        // Second row: cumulative product along axis 0
        // [1.0, 2.0, 1.0] then [1*3, 2*4, 1*2] = [3.0, 8.0, 2.0]
        let expected = vec![1.0, 2.0, 1.0, 3.0, 8.0, 2.0];
        assert_eq!(result.to_vec(), expected);
    }

    #[test]
    fn test_nancumprod_2d_axis1() {
        let a = Array::from_vec(vec![1.0, f64::NAN, 2.0, 3.0, 2.0, 4.0]).reshape(&[2, 3]);
        let result = nancumprod(&a, Some(1)).unwrap();
        assert_eq!(result.shape(), vec![2, 3]);
        // Row 0: [1.0, 1.0, 2.0] (NaN treated as 1)
        // Row 1: [3.0, 6.0, 24.0]
        let expected = vec![1.0, 1.0, 2.0, 3.0, 6.0, 24.0];
        assert_eq!(result.to_vec(), expected);
    }

    #[test]
    fn test_nancumprod_negative_axis() {
        let a = Array::from_vec(vec![1.0, f64::NAN, 2.0, 3.0]).reshape(&[2, 2]);
        let result = nancumprod(&a, Some(-1)).unwrap();
        assert_eq!(result.shape(), vec![2, 2]);
        // Row 0: [1.0, 1.0] (NaN treated as 1)
        // Row 1: [2.0, 6.0]
        let expected = vec![1.0, 1.0, 2.0, 6.0];
        assert_eq!(result.to_vec(), expected);
    }

    #[test]
    fn test_nanargmax_all_nan() {
        // When all values are NaN, should return 0
        let a = Array::from_vec(vec![f64::NAN, f64::NAN, f64::NAN]);
        let result = nanargmax(&a, None, false).unwrap();
        assert_eq!(result.to_vec(), vec![0]); // Returns 0 when all are NaN
    }

    #[test]
    fn test_nanargmin_all_nan() {
        // When all values are NaN, should return 0
        let a = Array::from_vec(vec![f64::NAN, f64::NAN, f64::NAN]);
        let result = nanargmin(&a, None, false).unwrap();
        assert_eq!(result.to_vec(), vec![0]); // Returns 0 when all are NaN
    }

    #[test]
    fn test_nancumprod_all_nan() {
        // When all values are NaN, result should be all ones
        let a = Array::from_vec(vec![f64::NAN, f64::NAN, f64::NAN]);
        let result = nancumprod(&a, None).unwrap();
        assert_eq!(result.to_vec(), vec![1.0, 1.0, 1.0]);
    }
}

#[cfg(test)]
mod test_histogram_bin_edges {
    use super::*;
    use numrs2::stats::{histogram_bin_edges, BinSpec};

    #[test]
    fn test_histogram_bin_edges_fixed_bins() {
        let data = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0]);
        let edges = histogram_bin_edges(&data, BinSpec::Count(5), None).unwrap();
        assert_eq!(edges.size(), 6); // 5 bins = 6 edges
        assert_relative_eq!(edges.get(&[0]).unwrap(), 1.0, epsilon = 1e-10);
        assert_relative_eq!(edges.get(&[5]).unwrap(), 5.0, epsilon = 1e-10);
    }

    #[test]
    fn test_histogram_bin_edges_with_range() {
        let data = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0]);
        let edges = histogram_bin_edges(&data, BinSpec::Count(4), Some((0.0, 8.0))).unwrap();
        assert_eq!(edges.size(), 5); // 4 bins = 5 edges
        assert_relative_eq!(edges.get(&[0]).unwrap(), 0.0, epsilon = 1e-10);
        assert_relative_eq!(edges.get(&[4]).unwrap(), 8.0, epsilon = 1e-10);
    }

    #[test]
    fn test_histogram_bin_edges_sqrt() {
        let data = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
        let edges = histogram_bin_edges(&data, BinSpec::Sqrt, None).unwrap();
        // sqrt(9) = 3, so 3 bins, 4 edges
        assert_eq!(edges.size(), 4);
    }

    #[test]
    fn test_histogram_bin_edges_sturges() {
        let data = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]);
        let edges = histogram_bin_edges(&data, BinSpec::Sturges, None).unwrap();
        // Sturges: ceil(log2(8) + 1) = ceil(3 + 1) = 4 bins, 5 edges
        assert_eq!(edges.size(), 5);
    }

    #[test]
    fn test_histogram_bin_edges_rice() {
        let data = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]);
        let edges = histogram_bin_edges(&data, BinSpec::Rice, None).unwrap();
        // Rice: ceil(2 * 8^(1/3)) = ceil(2 * 2) = 4 bins, 5 edges
        assert_eq!(edges.size(), 5);
    }

    #[test]
    fn test_histogram_bin_edges_auto() {
        let data = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]);
        let edges = histogram_bin_edges(&data, BinSpec::Auto, None).unwrap();
        // Auto should produce a reasonable number of bins
        assert!(edges.size() >= 2);
    }

    #[test]
    fn test_histogram_bin_edges_usize_into() {
        let data = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0]);
        let edges = histogram_bin_edges(&data, 5usize, None).unwrap();
        assert_eq!(edges.size(), 6);
    }

    #[test]
    fn test_histogram_bin_edges_string_into() {
        let data = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]);
        let edges = histogram_bin_edges(&data, "sqrt", None).unwrap();
        // sqrt(8) ≈ 2.83, ceil = 3 bins, 4 edges
        assert_eq!(edges.size(), 4);
    }

    #[test]
    fn test_histogram_bin_edges_empty_error() {
        let data: Array<f64> = Array::from_vec(vec![]);
        let result = histogram_bin_edges(&data, BinSpec::Count(5), None);
        assert!(result.is_err());
    }

    #[test]
    fn test_histogram_bin_edges_zero_bins_error() {
        let data = Array::from_vec(vec![1.0, 2.0, 3.0]);
        let result = histogram_bin_edges(&data, BinSpec::Count(0), None);
        assert!(result.is_err());
    }

    #[test]
    fn test_histogram_bin_edges_invalid_range_error() {
        let data = Array::from_vec(vec![1.0, 2.0, 3.0]);
        let result = histogram_bin_edges(&data, BinSpec::Count(5), Some((5.0, 1.0)));
        assert!(result.is_err());
    }

    #[test]
    fn test_histogram_bin_edges_uniform_spacing() {
        let data = Array::from_vec(vec![0.0, 10.0]);
        let edges = histogram_bin_edges(&data, BinSpec::Count(10), None).unwrap();
        assert_eq!(edges.size(), 11);

        // Check uniform spacing
        for i in 0..10 {
            let edge_i = edges.get(&[i]).unwrap();
            let edge_i_plus_1 = edges.get(&[i + 1]).unwrap();
            assert_relative_eq!(edge_i_plus_1 - edge_i, 1.0, epsilon = 1e-10);
        }
    }

    #[test]
    fn test_histogram_bin_edges_scott() {
        let data = Array::from_vec(vec![1.0, 2.0, 2.5, 3.0, 3.5, 4.0, 5.0, 6.0, 7.0, 8.0]);
        let edges = histogram_bin_edges(&data, BinSpec::Scott, None).unwrap();
        // Scott should produce a reasonable number of bins
        assert!(edges.size() >= 2);
    }

    #[test]
    fn test_histogram_bin_edges_fd() {
        let data = Array::from_vec(vec![1.0, 2.0, 2.5, 3.0, 3.5, 4.0, 5.0, 6.0, 7.0, 8.0]);
        let edges = histogram_bin_edges(&data, BinSpec::Fd, None).unwrap();
        // Freedman-Diaconis should produce a reasonable number of bins
        assert!(edges.size() >= 2);
    }

    #[test]
    fn test_histogram_bin_edges_doane() {
        let data = Array::from_vec(vec![1.0, 2.0, 2.5, 3.0, 3.5, 4.0, 5.0, 6.0, 7.0, 8.0]);
        let edges = histogram_bin_edges(&data, BinSpec::Doane, None).unwrap();
        // Doane should produce a reasonable number of bins
        assert!(edges.size() >= 2);
    }

    #[test]
    fn test_histogram_bin_edges_small_data() {
        let data = Array::from_vec(vec![1.0, 2.0]);
        let edges = histogram_bin_edges(&data, BinSpec::Auto, None).unwrap();
        // Should handle small data gracefully
        assert!(edges.size() >= 2);
    }
}

#[cfg(test)]
mod test_nan_axis_functions {
    use super::*;
    use numrs2::stats::{nanmax, nanmean, nanmin, nanprod, nanstd, nansum, nanvar};

    #[test]
    fn test_nanmean_2d_axis0() {
        // [[1.0, NaN, 3.0],
        //  [4.0, 5.0, 6.0]]
        let data = Array::from_vec(vec![1.0, f64::NAN, 3.0, 4.0, 5.0, 6.0]).reshape(&[2, 3]);
        let result = nanmean(&data, Some(0), false).unwrap();
        assert_eq!(result.shape(), vec![3]);
        let values = result.to_vec();
        assert_relative_eq!(values[0], 2.5, epsilon = 1e-10); // (1+4)/2
        assert_relative_eq!(values[1], 5.0, epsilon = 1e-10); // 5/1 (NaN ignored)
        assert_relative_eq!(values[2], 4.5, epsilon = 1e-10); // (3+6)/2
    }

    #[test]
    fn test_nanmean_2d_axis1() {
        let data = Array::from_vec(vec![1.0, f64::NAN, 3.0, 4.0, 5.0, 6.0]).reshape(&[2, 3]);
        let result = nanmean(&data, Some(1), false).unwrap();
        assert_eq!(result.shape(), vec![2]);
        let values = result.to_vec();
        assert_relative_eq!(values[0], 2.0, epsilon = 1e-10); // (1+3)/2
        assert_relative_eq!(values[1], 5.0, epsilon = 1e-10); // (4+5+6)/3
    }

    #[test]
    fn test_nanmean_2d_keepdims() {
        let data = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0]).reshape(&[2, 2]);
        let result = nanmean(&data, Some(0), true).unwrap();
        assert_eq!(result.shape(), vec![1, 2]);
    }

    #[test]
    fn test_nansum_2d_axis0() {
        let data = Array::from_vec(vec![1.0, f64::NAN, 3.0, 4.0, 5.0, 6.0]).reshape(&[2, 3]);
        let result = nansum(&data, Some(0), false).unwrap();
        assert_eq!(result.shape(), vec![3]);
        let values = result.to_vec();
        assert_relative_eq!(values[0], 5.0, epsilon = 1e-10); // 1+4
        assert_relative_eq!(values[1], 5.0, epsilon = 1e-10); // 5 (NaN ignored)
        assert_relative_eq!(values[2], 9.0, epsilon = 1e-10); // 3+6
    }

    #[test]
    fn test_nansum_2d_axis1() {
        let data = Array::from_vec(vec![1.0, f64::NAN, 3.0, 4.0, 5.0, 6.0]).reshape(&[2, 3]);
        let result = nansum(&data, Some(1), false).unwrap();
        assert_eq!(result.shape(), vec![2]);
        let values = result.to_vec();
        assert_relative_eq!(values[0], 4.0, epsilon = 1e-10); // 1+3
        assert_relative_eq!(values[1], 15.0, epsilon = 1e-10); // 4+5+6
    }

    #[test]
    fn test_nanmin_2d_axis0() {
        let data = Array::from_vec(vec![1.0, f64::NAN, 9.0, 4.0, 5.0, 6.0]).reshape(&[2, 3]);
        let result = nanmin(&data, Some(0), false).unwrap();
        assert_eq!(result.shape(), vec![3]);
        let values = result.to_vec();
        assert_relative_eq!(values[0], 1.0, epsilon = 1e-10); // min(1,4)
        assert_relative_eq!(values[1], 5.0, epsilon = 1e-10); // min(NaN,5) = 5
        assert_relative_eq!(values[2], 6.0, epsilon = 1e-10); // min(9,6)
    }

    #[test]
    fn test_nanmin_2d_axis1() {
        let data = Array::from_vec(vec![3.0, f64::NAN, 1.0, 4.0, 5.0, 6.0]).reshape(&[2, 3]);
        let result = nanmin(&data, Some(1), false).unwrap();
        assert_eq!(result.shape(), vec![2]);
        let values = result.to_vec();
        assert_relative_eq!(values[0], 1.0, epsilon = 1e-10); // min(3,NaN,1)
        assert_relative_eq!(values[1], 4.0, epsilon = 1e-10); // min(4,5,6)
    }

    #[test]
    fn test_nanmax_2d_axis0() {
        let data = Array::from_vec(vec![1.0, f64::NAN, 9.0, 4.0, 5.0, 6.0]).reshape(&[2, 3]);
        let result = nanmax(&data, Some(0), false).unwrap();
        assert_eq!(result.shape(), vec![3]);
        let values = result.to_vec();
        assert_relative_eq!(values[0], 4.0, epsilon = 1e-10); // max(1,4)
        assert_relative_eq!(values[1], 5.0, epsilon = 1e-10); // max(NaN,5) = 5
        assert_relative_eq!(values[2], 9.0, epsilon = 1e-10); // max(9,6)
    }

    #[test]
    fn test_nanmax_2d_axis1() {
        let data = Array::from_vec(vec![3.0, f64::NAN, 1.0, 4.0, 5.0, 6.0]).reshape(&[2, 3]);
        let result = nanmax(&data, Some(1), false).unwrap();
        assert_eq!(result.shape(), vec![2]);
        let values = result.to_vec();
        assert_relative_eq!(values[0], 3.0, epsilon = 1e-10); // max(3,NaN,1)
        assert_relative_eq!(values[1], 6.0, epsilon = 1e-10); // max(4,5,6)
    }

    #[test]
    fn test_nanprod_2d_axis0() {
        let data = Array::from_vec(vec![2.0, f64::NAN, 3.0, 4.0, 5.0, 6.0]).reshape(&[2, 3]);
        let result = nanprod(&data, Some(0), false).unwrap();
        assert_eq!(result.shape(), vec![3]);
        let values = result.to_vec();
        assert_relative_eq!(values[0], 8.0, epsilon = 1e-10); // 2*4
        assert_relative_eq!(values[1], 5.0, epsilon = 1e-10); // 5 (NaN ignored)
        assert_relative_eq!(values[2], 18.0, epsilon = 1e-10); // 3*6
    }

    #[test]
    fn test_nanvar_2d_axis0() {
        // [[1.0, 2.0],
        //  [3.0, 4.0]]
        let data = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0]).reshape(&[2, 2]);
        let result = nanvar(&data, Some(0), None, false).unwrap();
        assert_eq!(result.shape(), vec![2]);
        let values = result.to_vec();
        // Variance of [1,3]: mean=2, var = ((1-2)^2 + (3-2)^2)/2 = 1
        assert_relative_eq!(values[0], 1.0, epsilon = 1e-10);
        // Variance of [2,4]: mean=3, var = ((2-3)^2 + (4-3)^2)/2 = 1
        assert_relative_eq!(values[1], 1.0, epsilon = 1e-10);
    }

    #[test]
    fn test_nanstd_2d_axis0() {
        let data = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0]).reshape(&[2, 2]);
        let result = nanstd(&data, Some(0), None, false).unwrap();
        assert_eq!(result.shape(), vec![2]);
        let values = result.to_vec();
        // std = sqrt(var) = sqrt(1) = 1
        assert_relative_eq!(values[0], 1.0, epsilon = 1e-10);
        assert_relative_eq!(values[1], 1.0, epsilon = 1e-10);
    }

    #[test]
    fn test_nan_axis_all_nan_column() {
        // Test column with all NaN values
        let data = Array::from_vec(vec![f64::NAN, 2.0, f64::NAN, 4.0]).reshape(&[2, 2]);
        let result = nanmean(&data, Some(0), false).unwrap();
        let values = result.to_vec();
        assert!(values[0].is_nan()); // Column 0 is all NaN
        assert_relative_eq!(values[1], 3.0, epsilon = 1e-10); // (2+4)/2
    }

    #[test]
    fn test_nan_axis_out_of_bounds() {
        let data = Array::from_vec(vec![1.0, 2.0, 3.0, 4.0]).reshape(&[2, 2]);
        let result = nanmean(&data, Some(5), false);
        assert!(result.is_err());
    }
}