use num::traits::{Num, NumCast};

use super::error::DspError;

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

impl<'a, NumType> Upsampler<'a, NumType>
where
    NumType: Num + Copy,
{
    pub fn from<I>(iterator: I, factor: u16) -> Upsampler<'a, NumType>
    where
        I: Iterator<Item = Result<NumType, DspError>> + 'a,
    {
        Upsampler {
            factor: factor,
            state: 0,
            sample: Some(NumType::zero()),
            iterator: Box::new(iterator),
        }
    }
}

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

    fn next(&mut self) -> Option<Self::Item> {
        if self.state == 0 {
            self.sample = match self.iterator.next() {
                None => None,
                Some(Ok(x)) => Some(x),
                Some(Err(e)) => return Some(Err(e)),
            };
        }
        self.state = (self.state + 1) % self.factor;

        match self.sample {
            None => None,
            Some(x) => Some(Ok(x)),
        }
    }
}

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

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

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

    fn next(&mut self) -> Option<Self::Item> {
        let mut result = NumType::zero();
        for _ in 0..self.factor {
            match self.iterator.next() {
                Some(Ok(x)) => result = result + x,
                Some(Err(e)) => return Some(Err(e)),
                None => return None,
            }
        }
        result = result / NumType::from(self.factor).unwrap();

        Some(Ok(result))
    }
}

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

    #[test]
    fn upsampler_test() {
        let test_data = vec![1_f32, 2_f32, 3_f32, 4_f32, 5_f32];

        let upsampler = Upsampler::from(test_data.into_iter().map(|x| Ok(x)), 2);

        let result_data: Result<Vec<f32>, DspError> = upsampler.collect();
        assert_eq!(
            result_data,
            Ok(vec![
                1_f32, 1_f32, 2_f32, 2_f32, 3_f32, 3_f32, 4_f32, 4_f32, 5_f32, 5_f32
            ]),
        );
    }

    #[test]
    fn downsampler_test() {
        let test_data = vec![
            1_f32, 1_f32, 2_f32, 2_f32, 3_f32, 3_f32, 4_f32, 4_f32, 5_f32, 5_f32,
        ];

        let downsampler = Downsampler::from(test_data.into_iter().map(|x| Ok(x)), 2);

        let result_data: Result<Vec<f32>, DspError> = downsampler.collect();
        assert_eq!(result_data, Ok(vec![1_f32, 2_f32, 3_f32, 4_f32, 5_f32]));
    }
}