1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

use crate::{
    arange::ToArange,
    gridspace::{GridSpace, GridSpaceInterpolation},
    IntoGridSpace,
};
use array_bin_ops::Array;
use core::ops::Range;

/// [`Iterator`] returned by [`arange_grid`]
pub type ArangeGrid<T, const N: usize> = GridSpace<T, N>;

/// [`IntoIterator`] returned by [`ToGridSpace::into_grid_space`]
pub type IntoArangeGrid<T, const N: usize> = IntoGridSpace<T, N>;

/// Creates a grid space over the range made up of fixed step intervals
///
/// ```
/// use iter_num_tools::arange_grid;
///
/// let it = arange_grid([0.0, 0.0]..[1.0, 2.0], 0.5);
/// assert!(it.eq([
///     [0.0, 0.0], [0.5, 0.0],
///     [0.0, 0.5], [0.5, 0.5],
///     [0.0, 1.0], [0.5, 1.0],
///     [0.0, 1.5], [0.5, 1.5],
/// ]));
///
/// // different step count in each direction
/// let it = arange_grid([0.0, 0.0]..[1.0, 2.0], [0.5, 1.0]);
/// assert!(it.eq([
///     [0.0, 0.0], [0.5, 0.0],
///     [0.0, 1.0], [0.5, 1.0],
/// ]));
///
/// // even nd spaces
/// let it = arange_grid([0.0, 0.0, 0.0]..[2.0, 2.0, 2.0], 1.0);
/// assert!(it.eq([
///     [0.0, 0.0, 0.0], [1.0, 0.0, 0.0],
///     [0.0, 1.0, 0.0], [1.0, 1.0, 0.0],
///
///     [0.0, 0.0, 1.0], [1.0, 0.0, 1.0],
///     [0.0, 1.0, 1.0], [1.0, 1.0, 1.0],
/// ]));
/// ```
pub fn arange_grid<R, S, const N: usize>(range: R, step: S) -> ArangeGrid<R::Item, N>
where
    R: ToArangeGrid<S, N>,
{
    range.into_arange_grid(step).into_space()
}

/// Helper trait for [`arange_grid`]
pub trait ToArangeGrid<S, const N: usize> {
    /// The item that this is a arange grid over
    type Item;
    /// Create the arange grid
    fn into_arange_grid(self, step: S) -> IntoArangeGrid<Self::Item, N>;
}

impl<F: Copy, const N: usize> ToArangeGrid<[F; N], N> for Range<[F; N]>
where
    Range<F>: ToArange<F>,
{
    type Item = <Range<F> as ToArange<F>>::Item;

    fn into_arange_grid(self, step: [F; N]) -> IntoArangeGrid<Self::Item, N> {
        let Range { start, end } = self;

        let mut len = 1;
        let ranges = Array(start).zip_map(end, |start, end| start..end);
        let lerps = Array(ranges).zip_map(step, |range, step| {
            let space = range.into_arange(step);
            len *= space.len;
            space
        });

        IntoArangeGrid::new(len, GridSpaceInterpolation(lerps))
    }
}
impl<F: Copy, const N: usize> ToArangeGrid<F, N> for Range<[F; N]>
where
    Range<F>: ToArange<F>,
{
    type Item = <Range<F> as ToArange<F>>::Item;

    fn into_arange_grid(self, step: F) -> IntoArangeGrid<Self::Item, N> {
        let Range { start, end } = self;

        let mut len = 1;
        let lerps = Array(start).zip_map(end, |start, end| {
            let space = (start..end).into_arange(step);
            len *= space.len;
            space
        });

        IntoArangeGrid::new(len, GridSpaceInterpolation(lerps))
    }
}

#[cfg(test)]
mod tests {
    use crate::check_double_ended_iter;

    use super::*;

    #[test]
    fn test_arange_grid_exclusive() {
        check_double_ended_iter(
            arange_grid([0.0, 0.0]..[1.0, 2.0], [0.5, 1.0]),
            [[0.0, 0.0], [0.5, 0.0], [0.0, 1.0], [0.5, 1.0]],
        );
    }

    #[test]
    fn test_arange_grid_exclusive_len() {
        let mut it = arange_grid([0.0, 0.0]..[1.0, 2.0], 0.5);

        let mut expected_len = 8;

        assert_eq!(it.size_hint(), (expected_len, Some(expected_len)));

        while expected_len > 0 {
            assert_eq!(it.len(), expected_len);
            it.next();
            expected_len -= 1;
            assert_eq!(it.len(), expected_len);
            it.next_back();
            expected_len -= 1;
        }

        assert_eq!(it.len(), expected_len);
    }
}