Documentation
use crate::arr::core::{
    adaptative_result_shape, dim_expansion_for_shape, next_index, plain_index, shape_size,
};
use std::vec;

#[test]
fn shape_size_works() {
    let cases = [
        // shape, expected_size
        (vec![2, 6, 8], 2 * 6 * 8),
        (vec![1], 1),
        (vec![3, 2, 6, 8], 3 * 2 * 6 * 8),
    ];

    for (shape, expected_size) in cases.iter() {
        assert_eq!(*expected_size, shape_size(shape));
    }
}

#[test]
fn plain_index_works() {
    let cases = [
        // shape, index_shape, index, expected_plain_index
        (vec![2, 6, 1], vec![2, 6, 8], vec![0, 1, 3], 1 * 1),
        (
            vec![2, 1, 8],
            vec![2, 6, 8],
            vec![1, 3, 5],
            1 * (1 * 8) + 0 * 8 + 5,
        ),
    ];

    for (shape, index_shape, index, expected_plain_index) in cases.iter() {
        assert_eq!(
            *expected_plain_index,
            plain_index(shape, index_shape, index)
        );
    }
}

#[test]
fn next_index_works() {
    let shape = vec![2, 1, 2];
    let mut index = vec![0, 0, 0];

    let end = next_index(&shape, &mut index);
    assert_eq!(vec![0, 0, 1], index);
    assert!(!end);

    let end = next_index(&shape, &mut index);
    assert_eq!(vec![1, 0, 0], index);
    assert!(!end);

    let end = next_index(&shape, &mut index);
    assert_eq!(vec![1, 0, 1], index);
    assert!(!end);

    let end = next_index(&shape, &mut index);
    assert_eq!(vec![0, 0, 0], index);
    assert!(end);
}

#[test]
fn commutative_result_shape_works() {
    let cases = [
        // shape_a, shape_b, expected_result_shape
        (vec![2, 6, 8], vec![1], vec![2, 6, 8]),
        (vec![1], vec![3, 5, 8], vec![3, 5, 8]),
        (vec![1, 1], vec![3, 5, 8], vec![3, 5, 8]),
        (vec![1, 8], vec![3, 5, 8], vec![3, 5, 8]),
        (vec![10, 10], vec![10], vec![10, 10]),
    ];

    for (shape_a, shape_b, expected_result_shape) in cases.iter() {
        assert_eq!(
            *expected_result_shape,
            adaptative_result_shape(shape_a, shape_b)
        );
    }
}

#[test]
fn dim_expansion_for_shape_works() {
    assert_eq!(vec![1, 1, 5], dim_expansion_for_shape(&vec![5], 3));
    assert_eq!(vec![1, 1, 3, 5], dim_expansion_for_shape(&vec![3, 5], 4));
    assert_eq!(
        vec![8, 1, 3, 5],
        dim_expansion_for_shape(&vec![8, 1, 3, 5], 4)
    );
}