lasprs 0.14.1

use std::process::Output;

use super::*;
use crate::config::*;
use anyhow::{Result, bail};
use fft_convolver::FFTConvolver;
use num::Complex;

#[cfg_attr(feature = "python-bindings", pyclass)]
#[cfg_attr(feature = "python-bindings", gen_stub_pyclass)]
/// # Implementation of a FIR filter.
pub struct FIR {
    blocksize: usize, // of data to be filtered
    conv: FFTConvolver<Flt>,
    coeffs: Vec<Flt>, // impulse response; value is read by Python code
}
impl Debug for FIR {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let coeff_str = if self.coeffs.len() < 10 {
            format!("{:?}", self.coeffs)
        } else {
            format!("{:?}, ... , ", &self.coeffs[..2]).replace("]", "")
                + &format!("{:?}", &self.coeffs[self.coeffs.len() - 2..]).replace("[", "")
        };
        f.debug_struct("FIR")
            .field("blocksize", &self.blocksize)
            // .field("conv", &self.conv)
            .field("coeffs", &coeff_str)
            .finish()
    }
}

#[cfg(feature = "python-bindings")]
#[cfg_attr(feature = "python-bindings", pymethods)]
impl FIR {
    #[new]
    /// Create new FIR filter. See [FIR::new]
    // new(coeffs: &Vec<Flt>, blocksize: usize) -> Result<Self>
    pub fn new_py(coeffs: PyArrayLike1<Flt>, blocksize: usize) -> PyResult<Self> {
        Ok(FIR::new(coeffs.as_slice()?, blocksize)?)
    }

    /// Print FIR filter information in Python
    fn __repr__(&self) -> String {
        format!("{:#?}", self)
    }

    /// See: [FIR::filter()]
    #[pyo3(name = "filter")]
    pub fn filter_py<'py>(
        &mut self,
        py: Python<'py>,
        input: PyArrayLike1<Flt>,
    ) -> Result<PyArr1Flt<'py>> {
        let mut output = vec![0.0; input.len()?];
        self.filter(input.as_slice()?, &mut output);
        Ok(output.into_pyarray(py))
    }
}

impl FIR {
    /// Create new FIR filter
    ///
    /// # Args
    ///
    /// - coeffs: FIR filter coefficients
    /// - blocksize: blocksize at which the input data will be fed
    ///
    pub fn new(coeffs: &[Flt], blocksize: usize) -> Result<Self> {
        let mut conv = FFTConvolver::default();
        conv.init(blocksize, coeffs)?;
        Ok(Self {
            blocksize,
            conv,
            coeffs: coeffs.to_vec(),
        })
    }
}

impl FilterMethods for FIR {
    /// Filter data.
    ///
    /// # Args
    ///
    /// - input: Input data vector.
    ///
    /// # Panics
    ///
    /// Panics if the input vector length does not match the
    /// 'blocksize' given in FIR::new().
    ///
    fn filter(&mut self, input: &[Flt], output: &mut [Flt]) {
        debug_assert_eq!(
            input.len(),
            self.blocksize,
            "Input data does not match blocksize"
        );
        match self.conv.process(input, output) {
            Ok(_) => {}
            Err(e) => panic!("Error: {e}"),
        }
    }

    fn reset(&mut self) {
        self.conv.reset();
    }
}

#[cfg(test)]
mod test {
    use std::iter::Zip;
    use std::iter::{self, Repeat, repeat};

    use approx::assert_abs_diff_eq;
    use num::{complex::ComplexFloat, integer::sqrt, range};

    use super::*;

    /// Apply filter 'coeffs' to 'input' and check that the result is 'good'.
    fn helper_calc(input: Vec<Flt>, coeffs: Vec<Flt>, good: Vec<Flt>) {
        let mut firfilter = FIR::new(&coeffs, input.len()).unwrap();
        let mut calculated = vec![0.; input.len()];
        firfilter.filter(&input, &mut calculated);
        println!("Expected: {:?}", good);
        println!("Filtered: {:?}", calculated);
        let mut it = good.iter().zip(calculated.iter());
        assert!(it.all(|(g, c)| (g - c).abs() <= Flt::EPSILON * 10.));
    }

    /// FIR filter coefficients: 0,
    #[test]
    fn test_zero() {
        let coeffs = vec![0.0];
        let input = vec![1.3, 1.5, -0.45, -6.0, 5.1, 0.0, -0.0, 9.1];
        let good = vec![0.0; input.len()];

        helper_calc(input, coeffs, good);
    }

    /// FIR filter coefficients: 1,
    #[test]
    fn test_unity() {
        let coeffs = vec![1.0];
        let input = vec![1.3, 1.5, -0.45, -6.0, 5.1, 0.0, -0.0, 9.1];
        let good = input.clone();

        helper_calc(input, coeffs, good);
    }

    /// FIR filter coefficients: 1, 0, 0
    #[test]
    fn test_unity_with_trailing_zeros() {
        let mut coeffs = vec![0.0; 3];
        coeffs[0] = 1.0;
        let input = vec![1.3, 1.5, -0.45, -6.0, 5.1, 0.0, -0.0, 9.1];
        let good = input.clone();

        helper_calc(input, coeffs, good);
    }

    /// FIR filter coefficients: 2.5, 0, ... , 0
    #[test]
    fn test_gain() {
        let mut coeffs = vec![0.0; 10];
        coeffs[0] = 2.5;
        let input = vec![1.3, 1.5, -0.45, -6.0, 5.1, 0.0, -0.0, 9.1];
        let good = vec![3.25, 3.75, -1.125, -15., 12.75, 0., -0., 22.75];

        helper_calc(input, coeffs, good);
    }

    /// FIR filter coefficients: 0, 1 -> delays by 1 sample
    #[test]
    fn test_delay() {
        let delay = 3;
        let mut coeffs = vec![0.0; delay + 1];
        coeffs[delay] = 1.0;
        let input = vec![1.3, 1.5, -0.45, -6.0, 5.1, 0.0, -0.0, 9.1];
        let mut good = input[..input.len() - delay].to_vec();
        for _ in 0..delay {
            good.insert(0, 0.0);
        }

        helper_calc(input, coeffs, good);
    }

    /// Filter longer than signal (and 1 sample delay)
    #[test]
    fn test_long_filter() {
        let mut coeffs = vec![0.0; 100];
        coeffs[1] = 1.0;
        coeffs[51] = 2.0; // outside of input, should not be used
        let input = vec![1.3, 1.5, -0.45, -6.0, 5.1, 0.0, -0.0, 9.1];
        let mut good = input[..input.len() - 1].to_vec();
        good.insert(0, 0.0);

        helper_calc(input, coeffs, good);
    }

    // Feed input data in multiple blocks
    #[test]
    fn test_multiple_blocks() {
        let blocksize = 4;
        let mut coeffs = vec![0.0; 2];
        coeffs[1] = 1.0;
        let input = [1.3, 1.5, -0.45, -6.0, 5.1, 0.0, -0.0, 9.1];
        let mut good = input[..input.len() - 1].to_vec();
        good.insert(0, 0.0);

        let mut firfilter = FIR::new(&coeffs, blocksize).unwrap();

        let mut calculated = vec![];
        let Nblocks = input.len() / blocksize;
        for i in range(0, Nblocks) {
            let block_in: Vec<Flt> = input[i * blocksize..(i + 1) * blocksize].to_vec();
            let mut block_out = vec![0.; blocksize];
            firfilter.filter(&block_in, &mut block_out);
            block_out.iter().for_each(|&x| calculated.push(x));
        }
        println!("Expected: {:?}", good);
        println!("Filtered: {:?}", calculated);
        let mut it = good.iter().zip(calculated.iter());
        assert!(it.all(|(g, c)| (g - c).abs() <= Flt::EPSILON * 10.));
    }

    /// Wrong answer
    #[test]
    #[should_panic]
    fn test_panic_incorrect_answer() {
        let coeffs = vec![100.0];
        let input = vec![1.3, 1.5, -0.45, -6.0, 5.1, 0.0, -0.0, 9.1];
        let good = input.clone();

        helper_calc(input, coeffs, good);
    }

    /// Wrong blocksize
    #[test]
    #[should_panic]
    fn test_panic_blocksize() {
        let coeffs = [100.0];
        let input = [1.3, 1.5, -0.45, -6.0, 5.1, 0.0, -0.0, 9.1];

        let mut firfilter = FIR::new(&coeffs, 3).unwrap();
        let mut output = vec![0.; input.len()];
        firfilter.filter(&input, &mut output);
    }
}