adaptif 0.1.0

#![allow(
    clippy::needless_pass_by_value,
    reason = "PyArrays must be passed by value"
)]
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;

use numpy::{PyArray1, PyReadonlyArray1};

use super::macros::generate_filter_bindings;

use crate::Error;
use crate::algorithms::Algorithm;
use crate::filters::{FilterBase, ProcessingMode};
use crate::filters::{LmsFilter, LmsOptions, NlmsFilter, NlmsOptions, RlsFilter, RlsOptions};
use crate::types::signals::{InputSignal, NoiseReference};

impl From<Error> for PyErr {
    fn from(value: Error) -> Self {
        match value {
            Error::EmptyInputArr
            | Error::WindowSizeZero
            | Error::BlockSizeZero
            | Error::NoiseRefTooShort { .. }
            | Error::NonPositiveStepSize
            | Error::InvalidForgettingFactorRange
            | Error::NonPositivePInitScale
            | Error::NonPositiveEpsilon
            | Error::WeightsDontMatchWindowSize { .. } => PyValueError::new_err(value.to_string()),
        }
    }
}

// Python floats are double precision, so we use f64 instead of a generic
impl TryFrom<&PyReadonlyArray1<'_, f64>> for InputSignal<f64> {
    type Error = PyErr;
    fn try_from(value: &PyReadonlyArray1<f64>) -> Result<Self, Self::Error> {
        // TODO: if we split off a no_std core module, decide whether we want to copy
        // the input or use references (if we copy the input can be non-contiguous,
        // but it's more costly)

        let input_signal = value.as_array().iter().copied().collect::<Vec<f64>>();

        let input_signal = InputSignal::new(input_signal).map_err(PyErr::from)?;
        Ok(input_signal)
    }
}

impl TryFrom<&PyReadonlyArray1<'_, f64>> for NoiseReference<f64> {
    type Error = PyErr;
    fn try_from(value: &PyReadonlyArray1<f64>) -> Result<Self, Self::Error> {
        let noise_ref = value.as_array().iter().copied().collect::<Vec<f64>>();

        let noise_ref = NoiseReference::new(noise_ref).map_err(PyErr::from)?;
        Ok(noise_ref)
    }
}

#[derive(Debug, Clone, Copy)]
enum FilterOperation {
    Adapt,
    Filter,
}

// Because the wrappers for adapt() and filter() would only differ in one line,
// we use this underlying implementation.
fn adapt_filter_impl<'py, A>(
    filter: &mut FilterBase<f64, A>,
    py: Python<'py>,
    input_signal: PyReadonlyArray1<f64>,
    noise_ref: PyReadonlyArray1<f64>,
    op: FilterOperation,
) -> PyResult<Bound<'py, PyArray1<f64>>>
where
    A: Algorithm<f64>,
{
    let input_signal = InputSignal::try_from(&input_signal)?;
    let noise_ref = NoiseReference::try_from(&noise_ref)?;

    let output_signal = match op {
        FilterOperation::Adapt => filter.adapt(&input_signal, &noise_ref),
        FilterOperation::Filter => filter.filter(&input_signal, &noise_ref),
    }
    .map_err(PyErr::from)?;

    Ok(PyArray1::from_vec(py, output_signal))
}

#[pymodule]
mod adaptif {
    #[pymodule_export]
    use super::{LMSFilter, NLMSFilter, RLSFilter};
}

/// Least Mean Squares (LMS) adaptive filter.
///
/// Args:
///     `step_size`: Step size. Must be greater than 0.0.
///         Defaults to 0.1.
///     `window_size`: Number of filter weights. Must be greater than 0.
///         Defaults to 16.
///     `block_size`: Number of samples processed per block. Must be greater than 0.
///         Defaults to 1.
#[pyclass]
pub struct LMSFilter {
    inner: LmsFilter<f64>,
}
#[pymethods]
impl LMSFilter {
    #[new]
    #[pyo3(signature = (step_size=0.1, window_size=16, block_size=1))]
    fn new(step_size: f64, window_size: usize, block_size: usize) -> PyResult<Self> {
        let filter = LmsFilter::new(
            LmsOptions {
                step_size,
                window_size,
            },
            ProcessingMode::Block { block_size }, // Note that block_size: 1 is functionally
                                                  // equivalent to ProcessingMode::Sample
        )
        .map_err(PyErr::from)?;

        Ok(Self { inner: filter })
    }
}

generate_filter_bindings!(LMSFilter);

/// Normalized Least Mean Squares (NLMS) adaptive filter.
///
/// Args:
///     `step_size`: Step size. Must be greater than 0.0.
///         Defaults to 0.1.
///     `eps`: Normalization epsilon. Must be greater than 0.0
///         Defaults to 1e-8.
///     `window_size`: Number of filter weights. Must be greater than 0.
///         Defaults to 16.
///     `block_size`: Number of samples processed per block. Must be greater than 0.
///         Defaults to 1.
#[pyclass]
pub struct NLMSFilter {
    inner: NlmsFilter<f64>,
}
#[pymethods]
impl NLMSFilter {
    #[new]
    #[pyo3(signature = (step_size=0.01, eps=1e-8, window_size=16, block_size=1))]
    fn new(step_size: f64, eps: f64, window_size: usize, block_size: usize) -> PyResult<Self> {
        let filter = NlmsFilter::new(
            NlmsOptions {
                step_size,
                eps,
                window_size,
            },
            ProcessingMode::Block { block_size },
        )
        .map_err(PyErr::from)?;

        Ok(Self { inner: filter })
    }
}

generate_filter_bindings!(NLMSFilter);

/// Recursive Least Squares (RLS) adaptive filter.
///
/// Args:
///     `forgetting_factor`: RLS Forgetting factor. Must be greater than 0.0 and less-than or equal to 1.0.
///         Defaults to 0.999.
///     `p_init_scale`: Scalar value used in initializing the inverse correlation matrix. Must be greater than 0.0.
///         Defaults to 0.001.
///     `window_size`: Number of filter weights. Must be greater than 0.
///         Defaults to 8.
///     `block_size`: Number of samples processed per block. Must be greater than 0.
///         Defaults to 1.
#[pyclass]
pub struct RLSFilter {
    inner: RlsFilter<f64>,
}
#[pymethods]
impl RLSFilter {
    #[new]
    #[pyo3(signature = (forgetting_factor=0.999, p_init_scale=0.001, window_size=8, block_size=1))]
    fn new(
        forgetting_factor: f64,
        p_init_scale: f64,
        window_size: usize,
        block_size: usize,
    ) -> PyResult<Self> {
        let filter = RlsFilter::new(
            RlsOptions {
                forgetting_factor,
                p_init_scale,
                window_size,
            },
            ProcessingMode::Block { block_size },
        )
        .map_err(PyErr::from)?;

        Ok(Self { inner: filter })
    }
}

generate_filter_bindings!(RLSFilter);