ebu-plugs/dsp/src/biquad.rs

use std::{
    f64::{self, consts::SQRT_2},
    fmt::Display,
    sync::Arc,
};

use crate::{
    decibel::Decibel, SampleRate,
    traits::{IntFormatter, Processor},
};

#[derive(Default, Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum Slope {
    #[default]
    DB12,
    DB24,
    DB36,
    DB48,
}

#[derive(Default, Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum FilterMode {
    Peak,
    HighPass,
    HighShelf,
    LowPass,
    LowShelf,
    AllPass,
    BandPass,
    BandReject,
    #[default]
    None,
}
impl TryFrom<i32> for FilterMode {
    type Error = ();
    fn try_from(value: i32) -> Result<Self, Self::Error> {
        match value {
            0 => Ok(FilterMode::Peak),
            1 => Ok(FilterMode::HighPass),
            2 => Ok(FilterMode::HighShelf),
            3 => Ok(FilterMode::LowPass),
            4 => Ok(FilterMode::LowShelf),
            5 => Ok(FilterMode::AllPass),
            6 => Ok(FilterMode::BandPass),
            7 => Ok(FilterMode::BandReject),
            _ => Err(()),
        }
    }
}
impl From<FilterMode> for i32 {
    fn from(value: FilterMode) -> Self {
        match value {
            FilterMode::Peak => 0,
            FilterMode::HighPass => 1,
            FilterMode::HighShelf => 2,
            FilterMode::LowPass => 3,
            FilterMode::LowShelf => 4,
            FilterMode::AllPass => 5,
            FilterMode::BandPass => 6,
            FilterMode::BandReject => 7,
            FilterMode::None => 8,
        }
    }
}
impl TryFrom<&str> for FilterMode {
    type Error = ();

    fn try_from(value: &str) -> Result<Self, Self::Error> {
        match value {
            "Peak" => Ok(FilterMode::Peak),
            "Highpass" => Ok(FilterMode::HighPass),
            "High Shelf" => Ok(FilterMode::HighShelf),
            "Lowpass" => Ok(FilterMode::LowPass),
            "Low Shelf" => Ok(FilterMode::LowShelf),
            "Allpass" => Ok(FilterMode::AllPass),
            "Bandpass" => Ok(FilterMode::BandPass),
            "Notch" => Ok(FilterMode::BandReject),
            _ => Err(()),
        }
    }
}
impl Display for FilterMode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            FilterMode::Peak => "Peak",
            FilterMode::HighPass => "Highpass",
            FilterMode::HighShelf => "High Shelf",
            FilterMode::LowPass => "Lowpass",
            FilterMode::LowShelf => "Low Shelf",
            FilterMode::AllPass => "Allpass",
            FilterMode::BandPass => "Bandpass",
            FilterMode::BandReject => "Notch",
            FilterMode::None => "None",
        })
    }
}

impl IntFormatter for FilterMode {
    fn v2s() -> Arc<dyn Fn(i32) -> String + Send + Sync>
    where
        Self: Sized,
    {
        Arc::new(move |value| FilterMode::try_from(value).unwrap_or_default().to_string())
    }

    fn s2v() -> Arc<dyn Fn(&str) -> Option<i32> + Send + Sync>
    where
        Self: Sized,
    {
        Arc::new(move |value| FilterMode::try_from(value).map(i32::from).ok())
    }
}

#[derive(Default, Debug, Clone, Copy)]
struct FilterCoefficients {
    a1: f64,
    a2: f64,
    b0: f64,
    b1: f64,
    b2: f64,
}
impl FilterCoefficients {
    fn calculate(
        &mut self,
        mode: FilterMode,
        sample_rate: SampleRate,
        freq: f64,
        q: f64,
        gain_db: Decibel,
    ) {
        let v = 10f64.powf(gain_db.value().abs() / 20.0);
        let k = (f64::consts::PI * freq / sample_rate.0).tan();

        match mode {
            FilterMode::Peak => {
                if gain_db.value() >= 0.0 {
                    let norm = 1.0 / (1.0 + 1.0 / q * k + k * k);
                    self.b0 = (1.0 + v / q * k + k * k) * norm;
                    self.b1 = 2.0 * (k * k - 1.0) * norm;
                    self.b2 = (1.0 - v / q * k + k * k) * norm;
                    self.a1 = self.b1;
                    self.a2 = (1.0 - 1.0 / q * k + k * k) * norm;
                } else {
                    let norm = 1.0 / (1.0 + v / q * k + k * k);
                    self.b0 = (1.0 + 1.0 / q * k + k * k) * norm;
                    self.b1 = 2.0 * (k * k - 1.0) * norm;
                    self.b2 = (1.0 - 1.0 / q * k + k * k) * norm;
                    self.a1 = self.b1;
                    self.a2 = (1.0 - v / q * k + k * k) * norm;
                }
            }
            FilterMode::HighPass => {
                let norm = 1.0 / (1.0 + k / q + k * k);
                self.b0 = 1.0 * norm;
                self.b1 = -2.0 * self.b0;
                self.b2 = self.b0;
                self.a1 = 2.0 * (k * k - 1.0) * norm;
                self.a2 = (1.0 - k / q + k * k) * norm;
            }
            FilterMode::HighShelf => {
                if gain_db.value() >= 0.0 {
                    let norm = 1.0 / (1.0 + SQRT_2 * k + k * k);
                    self.b0 = (v + (2.0 * v).sqrt() * k + k * k) * norm;
                    self.b1 = 2.0 * (k * k - v) * norm;
                    self.b2 = (v - (2.0 * v).sqrt() * k + k * k) * norm;
                    self.a1 = 2.0 * (k * k - 1.0) * norm;
                    self.a2 = (1.0 - SQRT_2 * k + k * k) * norm;
                } else {
                    let norm = 1.0 / (v + (2.0 * v).sqrt() * k + k * k);
                    self.b0 = (1.0 + SQRT_2 * k + k * k) * norm;
                    self.b1 = 2.0 * (k * k - 1.0) * norm;
                    self.b2 = (1.0 - SQRT_2 * k + k * k) * norm;
                    self.a1 = 2.0 * (k * k - v) * norm;
                    self.a2 = (v - (2.0 * v).sqrt() * k + k * k) * norm;
                }
            }
            FilterMode::LowPass => {
                let norm = 1.0 / (1.0 + k / q + k * k);
                self.b0 = k * k * norm;
                self.b1 = 2.0 * self.b0;
                self.b2 = self.b0;
                self.a1 = 2.0 * (k * k - 1.0) * norm;
                self.a2 = (1.0 - k / q + k * k) * norm;
            }
            FilterMode::LowShelf => {
                if gain_db.value() >= 0.0 {
                    let norm = 1.0 / (1.0 + SQRT_2 * k + k * k);
                    self.b0 = (1.0 + (2.0 * v).sqrt() * k + v * k * k) * norm;
                    self.b1 = 2.0 * (v * k * k - 1.0) * norm;
                    self.b2 = (1.0 - (2.0 * v).sqrt() * k + v * k * k) * norm;
                    self.a1 = 2.0 * (k * k - 1.0) * norm;
                    self.a2 = (1.0 - SQRT_2 * k + k * k) * norm;
                } else {
                    let norm = 1.0 / (1.0 + (2.0 * v).sqrt() * k + v * k * k);
                    self.b0 = (1.0 + SQRT_2 * k + k * k) * norm;
                    self.b1 = 2.0 * (k * k - 1.0) * norm;
                    self.b2 = (1.0 - SQRT_2 * k + k * k) * norm;
                    self.a1 = 2.0 * (v * k * k - 1.0) * norm;
                    self.a2 = (1.0 - (2.0 * v).sqrt() * k + v * k * k) * norm;
                }
            }
            FilterMode::AllPass => {
                let norm = 1.0 / (1.0 + k / q + k * k);
                self.b0 = (1.0 - k / q + k * k) * norm;
                self.b1 = 2.0 * (k * k - 1.0) * norm;
                self.b2 = 1.0;
                self.a1 = self.b1;
                self.a2 = self.b0;
            }
            FilterMode::BandPass => {
                let norm = 1.0 / (1.0 + k / q + k * k);
                self.b0 = k / q * norm;
                self.b1 = 0.0;
                self.b2 = -self.b0;
                self.a1 = 2.0 * (k * k - 1.0) * norm;
                self.a2 = (1.0 - k / q + k * k) * norm;
            }
            FilterMode::BandReject => {
                let norm = 1.0 / (1.0 + k / q + k * k);
                self.b0 = (1.0 + k * k) * norm;
                self.b1 = 2.0 * (k * k - 1.0) * norm;
                self.b2 = self.b0;
                self.a1 = self.b1;
                self.a2 = (1.0 - k / q + k * k) * norm;
            }
            FilterMode::None => {}
        }
    }
}

#[derive(Debug, Default, Clone, Copy)]
struct State {
    s1: f64,
    s2: f64,
}

#[derive(Debug, Clone, Copy)]
pub struct BiquadFilterState {
    pub mode: FilterMode,
    pub cutoff: f64,
    pub q: f64,
    pub gain_db: Decibel,
    pub slope: Slope,
}
impl Default for BiquadFilterState {
    fn default() -> Self {
        Self {
            mode: FilterMode::None,
            cutoff: 2000.0,
            q: 0.707,
            gain_db: Decibel::from(0.0),
            slope: Slope::DB12,
        }
    }
}

#[derive(Debug, Clone)]
pub struct BiquadFilter {
    pub state: BiquadFilterState,
    sample_rate: SampleRate,
    coefficients: FilterCoefficients,
    slope_states: [State; 4],
    dirty: bool,
}

impl BiquadFilter {
    pub fn new(sample_rate: SampleRate) -> Self {
        Self::from_state(sample_rate, BiquadFilterState::default())
    }
    pub fn from_state(sample_rate: SampleRate, state: BiquadFilterState) -> Self {
        Self {
            sample_rate,
            coefficients: FilterCoefficients::default(),
            slope_states: [State::default(); 4],
            dirty: true,
            state
        }
    }
    pub fn set_mode(&mut self, mode: FilterMode) {
        self.dirty = true;
        self.state.mode = mode;
    }
    pub fn set_cutoff(&mut self, cutoff: f64) {
        self.dirty = true;
        self.state.cutoff = cutoff;
    }
    pub fn set_slope(&mut self, slope: Slope) {
        self.dirty = true;
        self.state.slope = slope;
    }
    pub fn set_q(&mut self, q: f64) {
        self.dirty = true;
        self.state.q = q;
    }
    pub fn set_gain<T>(&mut self, gain_db: T)
    where
        T: Into<Decibel>,
    {
        self.dirty = true;
        self.state.gain_db = gain_db.into();
    }

    fn update_coefficients(&mut self) {
        if !self.dirty {
            return;
        }
        self.dirty = false;

        self.coefficients.calculate(
            self.state.mode,
            self.sample_rate,
            self.state.cutoff,
            self.state.q,
            self.state.gain_db,
        );
    }
}

impl Processor for BiquadFilter {
    fn process(&mut self, sample: f64) -> f64 {
        self.update_coefficients();

        let iterations = match self.state.slope {
            Slope::DB12 => 1,
            Slope::DB24 => 2,
            Slope::DB36 => 3,
            Slope::DB48 => 4,
        };
        let mut current = sample;

        for i in 0..iterations {
            let result = self.coefficients.b0 * current + self.slope_states[i].s1;
            let s1 = self.coefficients.b1 * current - self.coefficients.a1 * result
                + self.slope_states[i].s2;
            let s2 = self.coefficients.b2 * current - self.coefficients.a2 * result;

            self.slope_states[i].s1 = s1;
            self.slope_states[i].s2 = s2;
            current = result;
        }

        if current.is_nan() || current.is_infinite() {
            0.0
        } else {
            current
        }
    }
}