use either::Either;
use std::{iter::Rev, ops::RangeInclusive};

pub fn split_from<T>(range: RangeInclusive<T>, x: T) -> [impl Iterator<Item = T> + Clone; 2]
where
    T: num::Integer + Copy,
    RangeInclusive<T>: Iterator<Item = T> + DoubleEndedIterator,
    Rev<RangeInclusive<T>>: Iterator<Item = T>,
{
    let in_range = range.contains(&x);
    let (start, end) = (*range.start(), *range.end());

    // RangeInclusive (1..=0), has 0 elements
    let base = Either::Right(T::one()..=T::zero());
    [
        if in_range && x > start + T::one() {
            Either::Left((start..=(x - T::one())).rev())
        } else {
            base.clone()
        },
        if in_range && x + T::one() < end {
            Either::Right((x + T::one())..=end)
        } else {
            base
        },
    ]
}

pub fn diag_raw<T>(
    i_chains: [impl Iterator<Item = T> + Clone; 2],
    j_chains: [impl Iterator<Item = T> + Clone; 2],
) -> [impl Iterator<Item = (T, T)> + Clone; 4]
where
    T: num::Integer + Copy,
    RangeInclusive<T>: Iterator<Item = T> + DoubleEndedIterator,
    Rev<RangeInclusive<T>>: Iterator<Item = T>,
{
    [(0, 0), (1, 1), (1, 0), (0, 1)].map(move |(a, b)| i_chains[a].clone().zip(j_chains[b].clone()))
}

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

    fn diag_test_helper<T>(
        i: T,
        j: T,
        range_i: RangeInclusive<T>,
        range_j: RangeInclusive<T>,
    ) -> [impl Iterator<Item = (T, T)> + Clone; 4]
    where
        T: num::Integer + Copy,
        RangeInclusive<T>: Iterator<Item = T> + DoubleEndedIterator,
        Rev<RangeInclusive<T>>: Iterator<Item = T>,
    {
        diag_raw(split_from(range_i, i), split_from(range_j, j))
    }

    #[test]
    fn split_test() {
        assert_eq!(
            split_from(0..=6, 2).map(Iterator::collect::<Vec<usize>>),
            [vec![1, 0], vec![3, 4, 5, 6]]
        );

        assert_eq!(
            split_from(0..=6, 0).map(Iterator::collect::<Vec<usize>>),
            [vec![], vec![1, 2, 3, 4, 5, 6]]
        );

        assert_eq!(
            split_from(0..=6, 6).map(Iterator::collect::<Vec<usize>>),
            [vec![5, 4, 3, 2, 1, 0], vec![]]
        );

        // test out-of-bounds and also generics
        assert_eq!(
            split_from::<i16>(-1i16..=4i16, 10i16).map(Iterator::collect::<Vec<i16>>),
            [const { Vec::new() }; 2]
        );
    }

    #[test]
    fn diag_test() {
        assert_eq!(
            diag_test_helper(2, 3, 0..=7, 0..=7).map(Iterator::collect::<Vec<(usize, usize)>>),
            [
                vec![(1, 2), (0, 1)],
                vec![(3, 4), (4, 5), (5, 6), (6, 7)],
                vec![(3, 2), (4, 1), (5, 0)],
                vec![(1, 4), (0, 5)]
            ]
        );
    }
}