adaptif 0.1.0

use std::ops::{Deref, DerefMut};

use crate::types::buffers::NoiseWindowView;
use crate::types::{Float, WindowSize};
use crate::{Error, Result};

use crate::algorithms::Algorithm;

/// Recursive Least Squares (RLS) algorithm.
///
/// Updates the filter weights using the [recursive least squares] algorithm.
/// RLS recursively finds the coefficients that
/// minimize an exponentially weighted sum of
/// squared errors to adapt the filter weights.
///
/// Unlike [least mean squares], which uses stochastic gradient
/// descent, RLS uses a deterministic least-squares approach.
///
/// RLS has fast convergence, at the cost of high computational complexity.
/// For computationally constrained applications,
/// LMS-family algorithms may be preferable.
///
/// [recursive least squares]: https://en.wikipedia.org/wiki/Recursive_least_squares_filter
/// [least mean squares]: https://en.wikipedia.org/wiki/Least_mean_squares_filter
#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Rls<F: Float> {
    /// Forgetting factor, often denoted with $\lambda$, used for weight updates.
    forgetting_factor: F,
    /// Positive scalar value, often denoted with $\delta$, used for initalizing the inverse correlation matrix $P\[n\]$.
    p_init_scale: F,
    /// Window size of the filter.
    /// Determines the size of the inverse correlation matrix and the kalman gain vector.
    window_size: WindowSize,
    /// Inverse correlation matrix, referred to as $P\[n\]$, for RLS updates.
    /// Its dimensions are (M, M), where M is the filter's window size.
    inverse_corr_matrix: InverseCorrMatrix<F>,
    /// Kalman Gain vector used in updating the filter's weights.
    /// Its length is equal to the filter's window size.
    kalman_gain: KalmanGain<F>,
}

impl<F: Float> Rls<F> {
    /// # Errors
    ///
    /// Returns an error if forgetting factor <= 0.0 or > 1.0.
    /// Returns an error if init scale <= 0.0.
    pub fn new(forgetting_factor: F, p_init_scale: F, window_size: WindowSize) -> Result<Self> {
        if forgetting_factor <= F::zero() || forgetting_factor > F::one() {
            return Err(Error::InvalidForgettingFactorRange);
        }

        if p_init_scale <= F::zero() {
            return Err(Error::NonPositivePInitScale);
        }
        Ok(Rls {
            forgetting_factor,
            p_init_scale,
            window_size,
            inverse_corr_matrix: InverseCorrMatrix::new(window_size, p_init_scale),
            kalman_gain: KalmanGain::new(window_size),
        })
    }
}

#[allow(clippy::doc_paragraphs_missing_punctuation, reason = "maths")]
/// The weights are updated at every step according to:
///
/// $$\mathbf{w}(n) = \mathbf{w}(n-1) + \mathbf{k}(n)e(n)$$
///
/// The Kalman gain vector is calculated as:
///
/// $$\mathbf{k}(n) = \frac{\mathbf{P}(n-1)\mathbf{x}(n)}{\lambda + \mathbf{x}^{T}(n)\mathbf{P}(n-1)\mathbf{x}(n)}$$
///
/// The inverse correlation matrix is then updated as:
///
/// $$\mathbf{P}(n) = \frac{1}{\lambda}\left(\mathbf{P}(n-1) - \mathbf{k}(n)\mathbf{x}^{T}(n)\mathbf{P}(n-1)\right)$$
///
/// The inverse correlation matrix is initialized as:
///
/// $$\mathbf{P}(0) = \delta\mathbf{I}$$
///
/// Where:
/// - $\mathbf{w}(n)$ is the filter weight vector.
/// - $\mathbf{x}(n)$ is the reference input vector.
/// - $e(n)$ is the scalar error for the current sample.
/// - $\mathbf{k}(n)$ is the Kalman gain vector.
/// - $\mathbf{P}(n)$ is the inverse correlation matrix.
/// - $\lambda$ is the forgetting factor, with $0 < \lambda \leq 1$.
/// - $\delta$ is  a positive initialization scalar.
/// - $\mathbf{I}$ is the identity matrix.
impl<F: Float> Algorithm<F> for Rls<F> {
    fn update_step(&mut self, weights: &mut [F], error: F, noise_window: NoiseWindowView<F>) {
        self.kalman_gain.update(
            &self.inverse_corr_matrix,
            &noise_window,
            self.forgetting_factor,
        );

        for (w, k) in weights.iter_mut().zip(self.kalman_gain.iter().copied()) {
            *w += k * error;
        }

        self.inverse_corr_matrix
            .update(&self.kalman_gain, &noise_window, self.forgetting_factor);
    }

    fn window_size(&self) -> WindowSize {
        self.window_size
    }
}

/// M-dimensional vector of Kalman gains, where M is the filter's window size.
#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
struct KalmanGain<F: Float> {
    inner: Box<[F]>,
}
impl<F: Float> KalmanGain<F> {
    pub fn new(window_size: WindowSize) -> Self {
        KalmanGain {
            inner: vec![F::zero(); *window_size].into_boxed_slice(),
        }
    }

    /// Updates the Kalman gain vector used in the RLS weight update.
    ///
    /// Kalman gain vector update calculated as:
    ///
    /// $`k_n` = \frac{P_{n-1} `x_n`}
    /// {\lambda + `x_n^T` P_{n-1} `x_n`}$.
    ///
    /// Existing Kalman gain buffer gets reused to store the numerator
    /// value in update calculation to avoid reallocating a new vector.
    fn update(
        &mut self,
        inverse_corr_matrix: &InverseCorrMatrix<F>,
        noise_ref: &NoiseWindowView<F>,
        lambda: F,
    ) {
        // reusing kalman buffer for getting numerator to avoid clone of numerator
        for (k_i, row) in self
            .iter_mut()
            .zip(inverse_corr_matrix.chunks_exact(noise_ref.len()))
        {
            *k_i = row
                .iter()
                .copied()
                .zip(noise_ref.iter())
                .map(|(px, x)| px * x)
                .sum::<F>();
        }

        let denominator = lambda
            + noise_ref
                .iter()
                .zip(self.iter().copied())
                .map(|(noise, num)| noise * num)
                .sum::<F>();

        for k_i in self.iter_mut() {
            *k_i /= denominator;
        }
    }
}
impl<F: Float> Deref for KalmanGain<F> {
    type Target = [F];
    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<F: Float> DerefMut for KalmanGain<F> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

/// Inverse Correlation Matrix with shape M * M, where M is the filter's window size.
#[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
struct InverseCorrMatrix<F: Float> {
    inner: Box<[F]>,
}
impl<F: Float> InverseCorrMatrix<F> {
    /// Creates an M x M-dimensional Identity matrix multiplied by a positive scalar `p_init_scale`.
    /// I.e. `p_init_scale` is the resulting value along main diagonal and all other values are initialized
    /// to zero.
    ///
    /// Serves as inverse correlation matrix in RLS algorithm, where M is the filter's window size.
    pub fn new(window_size: WindowSize, p_init_scale: F) -> Self {
        let mut p = vec![F::zero(); (*window_size) * (*window_size)].into_boxed_slice();

        for i in 0..(*window_size) {
            if let Some(elem) = p.get_mut(i * (*window_size) + i) {
                *elem = p_init_scale;
            }
        }
        InverseCorrMatrix { inner: p }
    }

    /// Updates the inverse correlation matrix used in the RLS weight update.
    /// Utilizes current Kalman gain `k_n`.
    ///
    /// Inverse correlation matrix update calculated as:
    ///
    /// $`P_n` = \frac{1}{\lambda}
    /// \left(P_{n-1} - `k_n` `x_n^T` P_{n-1}\right)$.
    ///
    /// The matrix is stored as a flat buffer and updated in place.
    fn update(&mut self, kalman_gain: &KalmanGain<F>, noise_ref: &NoiseWindowView<F>, lambda: F) {
        // We calculate [x^T P] by computing the dot product of
        // noise_ref (x) with each column in inverse_corr_matrix (P),
        // then update the values for that column in place.
        for col in 0..noise_ref.len() {
            // Get the j-th column in P and take its dot product with x.
            let p_col = self.iter().copied().skip(col).step_by(noise_ref.len());
            let xt_p_col = noise_ref.iter().zip(p_col).map(|(x, p)| x * p).sum::<F>();

            // Update the values in that column in place
            for (row, k_i) in kalman_gain.iter().copied().enumerate() {
                // index is into a flat buffer, so row * n gives us the start of each row
                let index = row * noise_ref.len() + col;

                // matrix is allocated as M x M elements, so index is valid
                // since kalman gain and noise_ref.len() have length M.
                let p_i = self
                    .get_mut(index)
                    .expect("Inverse corr matrix dimensions must match the reference window size");
                *p_i = (*p_i - k_i * xt_p_col) / lambda;
            }
        }
    }
}
impl<F: Float> Deref for InverseCorrMatrix<F> {
    type Target = [F];
    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<F: Float> DerefMut for InverseCorrMatrix<F> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::indexing_slicing, reason = "Tests")]
mod tests {
    use super::*;
    use crate::{
        test_utils::all_approx_equal,
        types::buffers::NoiseBuffer,
        types::{BlockSize, FilterWeights, WindowSize},
    };

    #[test]
    fn init_p_matrix_works() {
        let n = WindowSize::new(3).unwrap();
        let p_matrix = InverseCorrMatrix::new(n, 10.0);

        let expected = [10.0, 0.0, 0.0, 0.0, 10.0, 0.0, 0.0, 0.0, 10.0];

        assert_eq!(*p_matrix, expected);
    }

    #[test]
    fn update_kalman_gain_works() {
        let mut k = KalmanGain {
            inner: vec![0.0; 2].into_boxed_slice(),
        };
        let p = InverseCorrMatrix {
            inner: vec![1.0, 0.0, 0.0, 1.0].into_boxed_slice(),
        };
        let x_n = NoiseBuffer::from_slice(
            &[1.0, 2.0],
            WindowSize::new(2).unwrap(),
            BlockSize::new(1).unwrap(),
        );

        k.update(&p, &x_n.window(0).unwrap(), 1.0);

        let expected = [1.0 / 6.0, 1.0 / 3.0];

        assert!(all_approx_equal(k.iter(), expected.iter()));
    }

    #[test]
    fn update_p_matrix_works() {
        let mut p = InverseCorrMatrix {
            inner: vec![1.0, 0.0, 0.0, 1.0].into_boxed_slice(),
        };
        let k = KalmanGain {
            inner: vec![1.0 / 6.0, 1.0 / 3.0].into_boxed_slice(),
        };
        let x_n = NoiseBuffer::from_slice(
            &[1.0, 2.0],
            WindowSize::new(2).unwrap(),
            BlockSize::new(1).unwrap(),
        );

        p.update(&k, &x_n.window(0).unwrap(), 1.0);

        let expected = InverseCorrMatrix {
            inner: vec![5.0 / 6.0, -1.0 / 3.0, -1.0 / 3.0, 1.0 / 3.0].into_boxed_slice(),
        };

        assert!(all_approx_equal(p.iter(), expected.iter()));
    }

    #[test]
    fn update_rls_1() {
        let window_size = WindowSize::new(2).unwrap();
        let block_size = BlockSize::new(1).unwrap();

        let mut rls = Rls::new(0.5, 1.0, window_size).unwrap();
        let e_n = 2.0;
        let x_n = NoiseBuffer::from_slice(&[1.0, -1.0], window_size, block_size);
        let expected = [0.8, -0.8];
        let mut weights = FilterWeights::new(window_size);

        rls.update_step(&mut weights, e_n, x_n.window(0).unwrap());

        assert!(all_approx_equal(weights.iter(), expected.iter()));
    }

    #[test]
    fn update_rls_2() {
        let window_size = WindowSize::new(2).unwrap();
        let block_size = BlockSize::new(1).unwrap();

        let mut rls = Rls::new(1.0, 1.0, window_size).unwrap();
        let e_n = 1.0;
        let x_n = NoiseBuffer::from_slice(&[5.0, 2.0], window_size, block_size);
        let expected = [5.0 / 30.0, 2.0 / 30.0];
        let mut weights = FilterWeights::new(WindowSize::new(2).unwrap());

        rls.update_step(&mut weights, e_n, x_n.window(0).unwrap());

        assert!(all_approx_equal(weights.iter(), expected.iter()));
    }

    #[test]
    fn forgetting_factor_range() {
        let p_init_scale = 1.0;
        let window_size = WindowSize::new(8).unwrap();
        Rls::new(0.01, p_init_scale, window_size).unwrap();
        Rls::new(1.0, p_init_scale, window_size).unwrap();

        assert!(matches!(
            Rls::new(0.0, p_init_scale, window_size),
            Err(Error::InvalidForgettingFactorRange)
        ));
        assert!(matches!(
            Rls::new(-1.0, p_init_scale, window_size),
            Err(Error::InvalidForgettingFactorRange)
        ));

        assert!(matches!(
            Rls::new(2.0, p_init_scale, window_size),
            Err(Error::InvalidForgettingFactorRange)
        ));
    }

    #[test]
    fn p_init_scale_range() {
        let good_init_scale = 100.0;
        let window_size = WindowSize::new(8).unwrap();
        Rls::new(0.01, good_init_scale, window_size).unwrap();

        assert!(matches!(
            Rls::new(0.5, 0.0, window_size),
            Err(Error::NonPositivePInitScale)
        ));
        assert!(matches!(
            Rls::new(0.5, -1.0, window_size),
            Err(Error::NonPositivePInitScale)
        ));
    }

    #[test]
    fn window_size_works() {
        let window_size = WindowSize::new(16).unwrap();

        let rls = Rls::new(0.01, 1.0, window_size).unwrap();

        assert_eq!(rls.window_size(), window_size);
    }
}