use num::traits::Num;

use super::error::DspError;

pub struct Multiply<'a, NumType>
where
    NumType: Num + Copy,
{
    buffer: Option<NumType>,
    iterator_a: Box<dyn Iterator<Item = Result<NumType, DspError>> + 'a>,
    iterator_b: Box<dyn Iterator<Item = Result<NumType, DspError>> + 'a>,
}

impl<'a, NumType> Multiply<'a, NumType>
where
    NumType: Num + Copy,
{
    // Use two different types, because map will make its closure part of the type
    pub fn from<IA, IB>(iterator_a: IA, iterator_b: IB) -> Multiply<'a, NumType>
    where
        IA: Iterator<Item = Result<NumType, DspError>> + 'a,
        IB: Iterator<Item = Result<NumType, DspError>> + 'a,
    {
        Multiply {
            buffer: None,
            iterator_a: Box::new(iterator_a),
            iterator_b: Box::new(iterator_b),
        }
    }
}

impl<'a, NumType> Iterator for Multiply<'a, NumType>
where
    NumType: Num + Copy,
{
    type Item = Result<NumType, DspError>;

    fn next(&mut self) -> Option<Self::Item> {
        let value_a = match self.buffer {
            None => match self.iterator_a.next() {
                Some(Ok(a)) => a,
                Some(Err(e)) => return Some(Err(e)),
                None => return None,
            },
            Some(a) => {
                self.buffer = None;
                a
            }
        };

        let value_b = match self.iterator_b.next() {
            Some(Ok(b)) => b,
            Some(Err(e)) => {
                // Can't compute anything yet, buffer value_a until iterator_b stops blocking.
                self.buffer = Some(value_a);
                return Some(Err(e));
            }
            None => return None,
        };

        Some(Ok(value_a * value_b))
    }
}

pub struct MultiplyConst<'a, NumType>
where
    NumType: Num + Copy,
{
    factor: NumType,
    iterator: Box<dyn Iterator<Item = Result<NumType, DspError>> + 'a>,
}

impl<'a, NumType> MultiplyConst<'a, NumType>
where
    NumType: Num + Copy,
{
    pub fn from<I>(iterator: I, factor: NumType) -> MultiplyConst<'a, NumType>
    where
        I: Iterator<Item = Result<NumType, DspError>> + 'a,
    {
        MultiplyConst {
            factor: factor,
            iterator: Box::new(iterator),
        }
    }
}

impl<'a, NumType> Iterator for MultiplyConst<'a, NumType>
where
    NumType: Num + Copy,
{
    type Item = Result<NumType, DspError>;

    fn next(&mut self) -> Option<Self::Item> {
        let value = match self.iterator.next() {
            Some(Ok(x)) => x,
            Some(Err(e)) => return Some(Err(e)),
            None => return None,
        };

        Some(Ok(value * self.factor))
    }
}

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

    #[test]
    fn multiply_test() {
        let test_a = vec![1_f32, 2_f32, 3_f32];
        let test_b = vec![4_f32, 5_f32, 6_f32, 7_f32];

        let iter_a = test_a.into_iter().map(|a| Ok(a));
        let iter_b = test_b.into_iter().map(|b| Ok(b));

        let result = Multiply::from(iter_a, iter_b);

        let result_data: Result<Vec<f32>, DspError> = result.collect();
        assert_eq!(result_data, Ok(vec![4_f32, 10_f32, 18_f32]));
    }

    #[test]
    fn multiply_blocking_test() {
        let test_a = vec![1_f32, 2_f32, 3_f32];
        let test_b = vec![
            Ok(4_f32),
            Err(DspError::WouldBlock),
            Ok(5_f32),
            Ok(6_f32),
            Ok(7_f32),
        ];

        let iter_a = test_a.into_iter().map(|a| Ok(a));
        let iter_b = test_b.into_iter();

        let mut mult = Multiply::from(iter_a, iter_b);

        assert_eq!(mult.next(), Some(Ok(4_f32)));
        assert_eq!(mult.next(), Some(Err(DspError::WouldBlock)));
        assert_eq!(mult.next(), Some(Ok(10_f32)));
        assert_eq!(mult.next(), Some(Ok(18_f32)));
        assert_eq!(mult.next(), None);
    }

    #[test]
    fn multiplyconst_test() {
        let test_a = vec![1_f32, 2_f32, 3_f32];

        let result = MultiplyConst::from(test_a.into_iter().map(|x| Ok(x)), 5_f32);

        let result_data: Result<Vec<f32>, DspError> = result.collect();
        assert_eq!(result_data, Ok(vec![5_f32, 10_f32, 15_f32]));
    }
}