use num::traits::Float;

use super::error::DspError;

pub struct SquaringAMDemodulator<'a, NumType>
where
    NumType: Float,
{
    iterator: Box<dyn Iterator<Item = Result<NumType, DspError>> + 'a>,
}

impl<'a, NumType> SquaringAMDemodulator<'a, NumType>
where
    NumType: Float,
{
    pub fn from<I>(iterator1: I) -> SquaringAMDemodulator<'a, NumType>
    where
        I: Iterator<Item = Result<NumType, DspError>> + 'a,
    {
        SquaringAMDemodulator {
            iterator: Box::new(iterator1),
        }
    }
}

impl<'a, NumType> Iterator for SquaringAMDemodulator<'a, NumType>
where
    NumType: Float,
{
    type Item = Result<NumType, DspError>;

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

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

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

        let demod = SquaringAMDemodulator::from(test_data.into_iter().map(|x| Ok(x)));

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