lasprs 0.2.1

//! # Filter implemententations for biquads, series of biquads and banks of series of biquads.
//!
//! Contains [Biquad], [SeriesBiquad], and [BiquadBank]. These are all constructs that work on
//! blocks of input data, and apply filters on it. Todo: implement FIR filter.
#![allow(non_snake_case)]
use super::config::*;
use anyhow::{bail, Result};
use numpy::ndarray::{ArrayD, ArrayViewD, ArrayViewMutD};
use numpy::{IntoPyArray, PyArray1, PyArrayDyn, PyArrayLike1, PyReadonlyArrayDyn};
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::{pymodule, types::PyModule, PyResult};
use rayon::prelude::*;

pub trait Filter: Send {
    //! The filter trait is implemented by Biquad, SeriesBiquad, and BiquadBank

    /// Filter input to generate output. A vector of output floats is generated with the same
    /// length as input.
    fn filter(&mut self, input: &[Flt]) -> Vd;
    /// Reset the filter state(s). In essence, this makes sure that all memory of the past is
    /// forgotten.
    fn reset(&mut self);

    /// Required method for cloning a BiquadBank, such that arbitrary filter types can be used as
    /// their 'channels'.
    fn clone_dyn(&self) -> Box<dyn Filter>;
}
impl Clone for Box<dyn Filter> {
    fn clone(&self) -> Self {
        self.clone_dyn()
    }
}

/// # A biquad is a second order recursive filter structure.
#[cfg_attr(feature = "extension-module", pyclass)]
#[derive(Clone, Copy, Debug)]
pub struct Biquad {
    // State parameters
    w1: Flt,
    w2: Flt,
    // Filter coefficients - forward
    b0: Flt,
    b1: Flt,
    b2: Flt,
    // Filter coefficients - recursive
    // a0: Flt, // a0 is assumed one, not used
    a1: Flt,
    a2: Flt,
}
#[cfg(feature = "extension-module")]
#[cfg_attr(feature = "extension-module", pymethods)]
impl Biquad {
    #[new]
    /// Create new biquad filter. See [Biquad::new()]
    ///
    pub fn new_py<'py>(coefs: PyReadonlyArrayDyn<Flt>) -> PyResult<Self> {
        Ok(Biquad::new(&coefs.as_slice()?)?)
    }
    #[pyo3(name = "unit")]
    #[staticmethod]
    /// See: [Biquad::unit()]
    pub fn unit_py() -> Biquad {
        Biquad::unit()
    }
    #[pyo3(name = "firstOrderHighPass")]
    #[staticmethod]
    /// See: [Biquad::firstOrderHighPass()]
    pub fn firstOrderHighPass_py(fs: Flt, cuton_Hz: Flt) -> PyResult<Biquad> {
        Ok(Biquad::firstOrderHighPass(fs, cuton_Hz)?)
    }
    #[pyo3(name = "filter")]
    /// See: [Biquad::filter()]
    pub fn filter_py<'py>(
        &mut self,
        py: Python<'py>,
        input: PyArrayLike1<Flt>,
    ) -> PyResult<&'py PyArray1<Flt>> {
        Ok(self.filter(input.as_slice()?).into_pyarray(py))
    }
}
impl Biquad {
    /// Create new biquad filter from given filter coeficients
    ///
    /// # Args
    ///
    /// - coefs: Filter coefficients.
    ///
    pub fn new(coefs: &[Flt]) -> Result<Self> {
        match coefs {
            [b0, b1, b2, a0, a1, a2] => {
                if *a0 != 1.0  {
                    bail!("Coefficient a0 should be equal to 1.0")
                }
                Ok(Biquad { w1: 0., w2: 0., b0: *b0, b1: *b1, b2: *b2, a1: *a1, a2: *a2})
            },
            _ => bail!("Could not initialize biquad. Please make sure that the coefficients contain 6 terms")
        }
    }

    /// Create unit impulse response biquad filter. Input = output
    fn unit() -> Biquad {
        let filter_coefs = &[1., 0., 0., 1., 0., 0.];
        Biquad::new(filter_coefs).unwrap()
    }

    /// Initialize biquad as first order high pass filter.
    ///
    ///
    /// * fs: Sampling frequency in \[Hz\]
    /// * cuton_Hz: -3 dB cut-on frequency in \[Hz\]
    ///
    pub fn firstOrderHighPass(fs: Flt, cuton_Hz: Flt) -> Result<Biquad> {
        if fs <= 0. {
            bail!("Invalid sampling frequency: {} [Hz]", fs);
        }
        if cuton_Hz <= 0. {
            bail!("Invalid cuton frequency: {} [Hz]", cuton_Hz);
        }
        if cuton_Hz >= 0.98 * fs / 2. {
            bail!(
                "Invalid cuton frequency. We limit this to 0.98* fs / 2. Given value {} [Hz]",
                cuton_Hz
            );
        }

        let tau: Flt = 1. / (2. * pi * cuton_Hz);
        let facnum = 2. * fs * tau / (1. + 2. * fs * tau);
        let facden = (1. - 2. * fs * tau) / (1. + 2. * fs * tau);

        let coefs = [
            facnum,  // b0
            -facnum, // b1
            0.,      // b2
            1.,      // a0
            facden,  // a1
            0.,      // a2
        ];

        Ok(Biquad::new(&coefs).unwrap())
    }
    fn filter_inout(&mut self, inout: &mut [Flt]) {
        for sample in 0..inout.len() {
            let w0 = inout[sample] - self.a1 * self.w1 - self.a2 * self.w2;
            let yn = self.b0 * w0 + self.b1 * self.w1 + self.b2 * self.w2;
            self.w2 = self.w1;
            self.w1 = w0;
            inout[sample] = yn;
        }
        // println!("{:?}", inout);
    }
}
impl Filter for Biquad {
    fn filter(&mut self, input: &[Flt]) -> Vec<Flt> {
        let mut out = input.to_vec();
        self.filter_inout(&mut out);
        // println!("{:?}", out);
        out
    }
    fn reset(&mut self) {
        self.w1 = 0.;
        self.w2 = 0.;
    }
    fn clone_dyn(&self) -> Box<dyn Filter> {
        Box::new(self.clone())
    }
}

/// Series of biquads that filter sequentially on an input signal
///
/// # Examples
///
/// See [tests]
/// ```
#[derive(Clone, Debug)]
#[cfg_attr(feature = "extension-module", pyclass)]
pub struct SeriesBiquad {
    biqs: Vec<Biquad>,
}

#[cfg(feature = "extension-module")]
#[cfg_attr(feature = "extension-module", pymethods)]
impl SeriesBiquad {
    #[new]
    /// Create new series filter set. See [SeriesBiquad::new()]
    ///
    pub fn new_py<'py>(coefs: PyReadonlyArrayDyn<Flt>) -> PyResult<Self> {
        Ok(SeriesBiquad::new(&coefs.as_slice()?)?)
    }
    #[pyo3(name = "unit")]
    #[staticmethod]
    /// See: [Biquad::unit()]
    pub fn unit_py() -> SeriesBiquad {
        SeriesBiquad::unit()
    }
    #[pyo3(name = "filter")]
    /// See: [SeriesBiquad::filter()]
    pub fn filter_py<'py>(
        &mut self,
        py: Python<'py>,
        input: PyArrayLike1<Flt>,
    ) -> PyResult<&'py PyArray1<Flt>> {
        Ok(self.filter(input.as_slice()?).into_pyarray(py))
    }
    #[pyo3(name = "reset")]
    /// See: [SeriesBiquad::reset()]
    pub fn reset_py(&mut self) {
        self.reset();
    }
}
impl SeriesBiquad {
    /// Create a new series biquad, having an arbitrary number of biquads.
    ///
    /// # Arguments
    ///
    /// * `filter_coefs` - Vector of biquad coefficients, stored in a single array. The first six
    /// for the first biquad, and so on.
    ///
    ///
    pub fn new(filter_coefs: &[Flt]) -> Result<SeriesBiquad> {
        if filter_coefs.len() % 6 != 0 {
            bail!(
                "filter_coefs should be multiple of 6, given: {}.",
                filter_coefs.len()
            );
        }
        let nfilters = filter_coefs.len() / 6;

        let mut biqs: Vec<Biquad> = Vec::with_capacity(nfilters);
        for coefs in filter_coefs.chunks(6) {
            let biq = Biquad::new(coefs)?;
            biqs.push(biq);
        }

        if biqs.len() == 0 {
            bail!("No filter coefficients given!");
        }

        Ok(SeriesBiquad { biqs })
    }

    /// Unit impulse response series biquad. Input = output
    pub fn unit() -> SeriesBiquad {
        let filter_coefs = &[1., 0., 0., 1., 0., 0.];
        SeriesBiquad::new(filter_coefs).unwrap()
    }
    fn clone_dyn(&self) -> Box<dyn Filter> {
        Box::new(self.clone())
    }
}

impl Filter for SeriesBiquad {
    //! Filter input by applying all biquad filters in series on each input sample, to obtain the
    //! output samples.
    //!
    fn filter(&mut self, input: &[Flt]) -> Vd {
        let mut inout = input.to_vec();
        for biq in self.biqs.iter_mut() {
            biq.filter_inout(&mut inout);
        }
        inout
    }
    fn reset(&mut self) {
        self.biqs.iter_mut().for_each(|f| f.reset());
    }
    fn clone_dyn(&self) -> Box<dyn Filter> {
        Box::new(self.clone())
    }
}

#[cfg_attr(feature = "extension-module", pyclass)]
#[derive(Clone)]
/// Multiple biquad filter that operate in parallel on a signal, and can apply a gain value to each
/// of the returned values. The BiquadBank can be used to decompose a signal by running it through
/// parallel filters, or it can directly be used to eq a signal. For the latter process, also a
/// gain can be applied when the output is made as the sum of the filtered inputs for each biquad.
///
/// # Detailed description
///
/// Below is an example of the signal flow is for the case of three SeriesBiquad filters, `h1`,
/// `h2` and `h3`:
///
/// ```markdown
///
///            analysis()       gain application             sum()              
///                                                               
///                  +------+   +-----+                    +------------+
///            +-----+ h1   |---+ g1  +-------------+      |            |
///            |     +------+   +-----+             ++     |   +------  |
///            |                                     +-----|   ++       |
///  Input     |     +------+   +-----+                    |    ++      |   Output of filter()
///--------|>--+-----+ h2   |---+ g2  |--------------------|     +-+    +----------------|>
///            |     ---+---+   +-----+                    |     +-+    |
///            |                                     +-----|    ++      |
///            |     +------+   +-----+             ++     |   ++       |
///            +-----+ h3   |---+ g3  |-------------+      |   +------  |
///                  +------+   +-----+                    |            |
///                           |                            +------------+
///                           |                  
///                           | Output of analysis() method (optional)
///                           +-------------|>
/// ```
pub struct BiquadBank {
    biqs: Vec<Box<dyn Filter>>,
    gains: Vec<Flt>,
}

#[cfg(feature = "extension-module")]
#[cfg_attr(feature = "extension-module", pymethods)]
/// Methods to wrap it in Python
impl BiquadBank {
    #[new]
    /// Create new biquadbank filter set. See [BiquadBank::new()]
    ///
    pub fn new_py<'py>(coefs: PyReadonlyArrayDyn<Flt>) -> PyResult<Self> {
        let mut filts = vec![];
        for col in coefs.as_array().columns() {
            match col.as_slice() {
                Some(colslice) => {
                    let new_ser = SeriesBiquad::new(colslice)?;
                    filts.push(new_ser.clone_dyn());
                }
                None => {
                    return Err(PyValueError::new_err("Error generating column"));
                }
            }
        }
        Ok(BiquadBank::new(filts))
    }
    #[pyo3(name = "filter")]
    /// See: [BiquadBank::filter()]
    pub fn filter_py<'py>(
        &mut self,
        py: Python<'py>,
        input: PyArrayLike1<Flt>,
    ) -> PyResult<&'py PyArray1<Flt>> {
        Ok(self.filter(input.as_slice()?).into_pyarray(py))
    }
    #[pyo3(name = "reset")]
    /// See: [BiquadBank::reset()]
    pub fn reset_py(&mut self) {
        self.reset();
    }

    #[pyo3(name = "set_gains")]
    /// See: [BiquadBank::set_gains()]
    pub fn set_gains_py<'py>(&mut self, py: Python<'py>, gains: PyArrayLike1<Flt>) -> PyResult<()> {
        if gains.len() != self.len() {
            return Err(PyValueError::new_err("Invalid number of provided gains"));
        }
        self.set_gains(gains.as_slice()?);
        Ok(())
    }
    #[pyo3(name = "set_gains_dB")]
    /// See: [BiquadBank::set_gains_dB()]
    pub fn set_gains_dB_py<'py>(&mut self, py: Python<'py>, gains_dB: PyArrayLike1<Flt>) -> PyResult<()> {
        if gains_dB.len() != self.len() {
            return Err(PyValueError::new_err("Invalid number of provided gains"));
        }
        self.set_gains_dB(gains_dB.as_slice()?);
        Ok(())

    }
    #[pyo3(name = "len")]
    /// See: [BiquadBank::len()]
    pub fn len_py(&self) -> usize {
        self.len()
    }
}

impl BiquadBank {
    /// Create new biquad bank. Initialized from given vector of series biquads.
    pub fn new(biqs: Vec<Box<dyn Filter>>) -> BiquadBank {
        let gains = vec![1.0; biqs.len()];
        BiquadBank { biqs, gains }
    }
    /// Return the number of parallel filters installed.
    pub fn len(&self) -> usize {
        self.biqs.len()
    }

    /// Set the gains for each of the biquad. The gains are not used in the analyisis phase, but in
    /// the reconstruction phase, so when BiquadBank::filter() is run on an input signal.
    ///
    /// # Panics
    ///
    /// When gains_dB.len() != to the number of filters.
    pub fn set_gains_dB(&mut self, gains_dB: &[Flt]) {
        if gains_dB.len() != self.gains.len() {
            panic!("Invalid gains size!");
        }
        self.gains
            .iter_mut()
            .zip(gains_dB)
            .for_each(|(g, gdB)| *g = Flt::powf(10., gdB / 20.));
    }
    /// Set linear gain values for each biquad. Same comments hold as for
    /// [BiquadBank::set_gains_dB()].
    pub fn set_gains(&mut self, gains: &[Flt]) {
        if gains.len() != self.gains.len() {
            panic!("Invalid gains size!");
        }
        // This could be done more efficient, but it does not matter. How often would you change
        // the gain values?
        self.gains.clone_from(&gains.to_vec());
    }
    /// Analysis step. Runs input signal through all filters. Outputs a vector of output results,
    /// one for each filter in the bank.

    pub fn analysis(&mut self, input: &[Flt]) -> Vec<Vd> {
        // Filtered output for each filter in biquad bank
        let filtered_out: Vec<Vd> = self
            .biqs
            .par_iter_mut()
            // .iter_mut()
            .map(|biq| biq.filter(input))
            .collect();
        filtered_out
    }
}

impl Filter for BiquadBank {
    fn filter(&mut self, input: &[Flt]) -> Vd {
        // Sum of filter output times gains
        let filtered_out = self.analysis(input);

        let mut out: Vd = vec![0.; input.len()];

        for (f, g) in filtered_out.iter().zip(&self.gains) {
            for (outi, fi) in out.iter_mut().zip(f) {
                // Sum and multiply by gain value
                *outi += g * fi;
            }
        }
        out
    }
    fn reset(&mut self) {
        self.biqs.iter_mut().for_each(|b| b.reset());
    }
    fn clone_dyn(&self) -> Box<dyn Filter> {
        Box::new(self.clone())
    }
}

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

    #[test]
    fn test_biquad1() {
        let mut ser = Biquad::unit();
        let inp = vec![1., 0., 0., 0., 0., 0.];
        let filtered = ser.filter(&inp);
        assert_eq!(&filtered, &inp);
    }
    #[test]
    #[should_panic]
    fn test_biquad2() {
        // A a0 coefficient not in the right place, meaning we panic on unwrap
        let filter_coefs = vec![1., 0., 0., 0., 0., 0.];
        let mut ser = SeriesBiquad::new(&filter_coefs).unwrap();
        let inp = vec![1., 0., 0., 0., 0., 0.];
        let filtered = ser.filter(&inp);
        assert_eq!(&filtered, &inp);
    }
    #[test]
    fn test_biquad3() {
        let filter_coefs = vec![0.5, 0.5, 0., 1., 0., 0.];
        let mut ser = SeriesBiquad::new(&filter_coefs).unwrap();

        let mut inp = vec![1., 0., 0., 0., 0., 0.];
        let filtered = ser.filter(&inp);

        // Change input to see match what should come out of output
        inp[0] = 0.5;
        inp[1] = 0.5;
        assert_eq!(&inp, &filtered);
    }
    #[test]
    fn test_biquadbank1() {
        //! Creates two unit filters with gains of 0.5. Runs the input signal through these filters
        //! in parallel and check if input == output.
        let ser = Biquad::unit();
        let mut biq = BiquadBank::new(vec![ser.clone_dyn(), ser.clone_dyn()]);
        biq.set_gains(&[0.5, 0.5]);
        let inp = vec![1., 0., 0., 0., 0., 0.];
        let out = biq.filter(&inp);
        assert_eq!(&out, &inp);
    }
}