lasprs 0.14.3

use anyhow::{Result, bail};

use super::FIR;
use crate::{
    Flt, StrictlyPositive, Vd,
    filter::{Filter, FilterMethods},
    twopi, *,
};

enum State {
    // Fir 0 active
    FIR0Active,
    // Fir 1 active
    FIR1Active,
    // boolean: true = switching from 0 to 1, false = switching from 1 to 0.
    // Index: number of samples processed so far.
    Switching(bool, usize),
}

/// Adaptable FIR filter that smoothly switches between two FIR filters when
/// updating the coefficients. The switching uses a skew sine for fading in/out.
#[cfg_attr(feature = "python-bindings", gen_stub_pyclass, pyclass)]
pub struct AdaptableFIR {
    // Two FIR filters to switch between.
    firs: [FIR; 2],
    // Current state of the filter, either FIR0Active, FIR1Active, or Switching.
    state: State,
    // Number of samples over which to transition between filters when updating
    // coefficients. This is calculated as `fs * transition_time`.
    transition_samples: usize,
    // Number of samples to process per block.
    block_size: usize,
}

impl AdaptableFIR {
    /// Creates a new `AdaptableFIR` filter.
    ///
    /// # Arguments
    ///
    /// * `fs` - Sampling frequency.
    /// * `block_size` - Number of samples to process per block.
    /// * `transition_time` - Time scale for transitioning between filters when updating the coefficients.
    /// * `init_coeffs` - Initial filter coefficients. Optionally. if None, defaults to a zero-coefficient filter.
    pub fn new(
        fs: StrictlyPositive,
        block_size: usize,
        transition_time: StrictlyPositive,
        init_coeffs: Option<&[Flt]>,
    ) -> Self {
        let coefs = init_coeffs.unwrap_or(&[0.]);
        let transition_samples = (*fs * *transition_time) as usize;
        assert!(transition_samples > 0);

        Self {
            block_size,
            transition_samples,
            firs: [
                FIR::new(coefs, block_size).unwrap(),
                FIR::new(coefs, block_size).unwrap(),
            ],
            state: State::FIR0Active,
        }
    }

    /// Updates the filter coefficients, start switching between the two FIR filters.
    ///
    /// # Arguments
    ///
    /// * `coefs` - The new filter coefficients to set.
    pub fn updateCoefficients(&mut self, coefs: &[Flt]) -> Result<()> {
        match self.state {
            State::FIR0Active => {
                self.firs[1] = FIR::new(coefs, self.block_size).unwrap();
                self.state = State::Switching(true, 0);
                Ok(())
            }
            State::FIR1Active => {
                self.firs[0] = FIR::new(coefs, self.block_size).unwrap();
                self.state = State::Switching(false, 0);
                Ok(())
            }
            State::Switching(..) => {
                bail!("Cannot update coefficients while an update is in progress");
            }
        }
    }
}
impl FilterMethods for AdaptableFIR {
    fn filter(&mut self, input: &[Flt], output: &mut [Flt]) {
        let Self {
            firs,
            state,
            transition_samples,
            ..
        } = self;
        match state {
            State::FIR0Active => firs[0].filter(input, output),
            State::FIR1Active => firs[1].filter(input, output),
            State::Switching(from_0to1, n) => {
                firs[0].filter(input, output);
                let mut o1 = vec![0.0; output.len()];
                firs[1].filter(input, &mut o1);
                for (o0i, o1i) in output.iter_mut().zip(o1.iter()) {
                    let (gain0, gain1) = calc_gains(*from_0to1, *n, *transition_samples);
                    *o0i = gain0 * *o0i + gain1 * o1i;
                    *n += 1;
                }
                if n >= transition_samples {
                    if *from_0to1 {
                        *state = State::FIR1Active;
                    } else {
                        *state = State::FIR0Active;
                    }
                }
            }
        }
    }

    fn reset(&mut self) {
        match self.state {
            State::FIR0Active => self.firs[0].reset(),
            State::FIR1Active => self.firs[1].reset(),
            State::Switching(_, _) => {
                self.firs[0].reset();
                self.firs[1].reset();
            }
        }
    }
}

#[cfg(feature = "python-bindings")]
#[cfg_attr(feature = "python-bindings", gen_stub_pymethods, pymethods)]
impl AdaptableFIR {
    #[new]
    fn new_py(fs: StrictlyPositive, block_size: usize, transition_time: StrictlyPositive) -> Self {
        Self::new(fs, block_size, transition_time, None)
    }
    #[pyo3(name = "updateCoefficients")]
    fn updateCoefficients_py(&mut self, coefs: PyReadonlyArray1<Flt>) -> Result<()> {
        let coefs = coefs.as_slice().unwrap();
        self.updateCoefficients(coefs)
    }
    #[pyo3(name = "filter")]
    fn filter_py<'py>(
        &mut self,
        py: Python<'py>,
        input: PyReadonlyArray1<Flt>,
    ) -> PyResult<Bound<'py, PyArray1<Flt>>> {
        let mut output = vec![0.0; input.len()?];
        match input.as_slice().ok() {
            Some(i) => {
                self.filter(i, &mut output);
                Ok(output.into_pyarray(py))
            }
            None => {
                // Create a temporary vector from the input array and filter it
                let input = &input.as_array().iter().copied().collect::<Vec<_>>();
                self.filter(input, &mut output);
                Ok(output.into_pyarray(py))
            }
        }
    }
}

// Calculate gains when mixing two channels. Uses a skewed sine for fading
// between them.
//
// # Returns
//
// A tuple of `(gain0, gain1)` where `gain0` is the gain for the first channel
// and `gain1` is the gain for the second channel.
#[inline]
fn calc_gains(from_0to1: bool, n: usize, transition_samples: usize) -> (Flt, Flt) {
    let gainA = skewsine((n as Flt) / (transition_samples - 1) as Flt);
    let gainB = 1. - gainA;
    if from_0to1 {
        (gainB, gainA)
    } else {
        (gainA, gainB)
    }
}

/// Skewed sine function used for fading between two channels.
#[inline]
fn skewsine(val: Flt) -> Flt {
    if val < 0. {
        return 0.;
    }
    if val > 1. {
        return 1.;
    }
    val - 1. / twopi * Flt::sin(twopi * val)
}