oximedia-effects 0.1.8

//! Psychoacoustic bass enhancer.
//!
//! Generates harmonics of bass frequencies to improve the perceived low-end
//! on small speakers and headphones that cannot reproduce sub-bass content
//! directly.  The technique exploits the "missing fundamental" psychoacoustic
//! phenomenon: the brain reconstructs the fundamental pitch from its overtones.
//!
//! # Algorithm
//!
//! 1. A second-order Butterworth **lowpass crossover** at `frequency_hz` isolates
//!    the bass band.
//! 2. The bass band is passed through a **waveshaping harmonic exciter** that
//!    introduces harmonics via a `tanh`-based soft-clipper.  The drive parameter
//!    controls how many harmonics are generated and their relative level.
//! 3. The high-frequency components generated by the waveshaper are extracted
//!    with a **highpass filter** tuned to the crossover frequency, preventing
//!    the original bass from being added twice.
//! 4. The harmonic content is **mixed back** with the full original signal.
//!
//! # Example
//!
//! ```
//! use oximedia_effects::bass_enhancer::{BassEnhancer, BassEnhancerConfig};
//!
//! let config = BassEnhancerConfig::default();
//! let mut enhancer = BassEnhancer::new(config, 48_000.0);
//!
//! let mut buf = vec![0.3f32; 512];
//! enhancer.process(&mut buf);
//! ```

#![allow(clippy::cast_precision_loss)]

use std::f32::consts::PI;

// ---------------------------------------------------------------------------
// Configuration
// ---------------------------------------------------------------------------

/// Configuration for the psychoacoustic bass enhancer.
#[derive(Debug, Clone, PartialEq)]
pub struct BassEnhancerConfig {
    /// Crossover frequency in Hz.  Frequencies below this are treated as
    /// "bass" and enhanced.  Typical: 60–120 Hz.
    pub frequency_hz: f32,
    /// Number of harmonics to generate (1–8).  Each harmonic is the
    /// `n`-th overtone of the bass content.
    pub harmonics: u32,
    /// Drive level for the waveshaper (0.0–1.0).  Higher = more harmonic
    /// distortion = stronger enhancement.
    pub drive: f32,
    /// Mix level of the harmonic content into the output (0.0–1.0).
    pub mix: f32,
}

impl Default for BassEnhancerConfig {
    fn default() -> Self {
        Self {
            frequency_hz: 80.0,
            harmonics: 3,
            drive: 0.5,
            mix: 0.4,
        }
    }
}

impl BassEnhancerConfig {
    /// Gentle enhancement preset.
    #[must_use]
    pub fn gentle() -> Self {
        Self {
            frequency_hz: 100.0,
            harmonics: 2,
            drive: 0.25,
            mix: 0.25,
        }
    }

    /// Aggressive enhancement preset.
    #[must_use]
    pub fn aggressive() -> Self {
        Self {
            frequency_hz: 80.0,
            harmonics: 5,
            drive: 0.8,
            mix: 0.6,
        }
    }

    /// Validate parameters.
    pub fn validate(&self) -> Result<(), String> {
        if self.frequency_hz <= 0.0 {
            return Err(format!(
                "frequency_hz must be positive, got {}",
                self.frequency_hz
            ));
        }
        if self.harmonics == 0 || self.harmonics > 16 {
            return Err(format!(
                "harmonics must be between 1 and 16, got {}",
                self.harmonics
            ));
        }
        if !(0.0..=1.0).contains(&self.drive) {
            return Err(format!("drive must be in [0.0, 1.0], got {}", self.drive));
        }
        if !(0.0..=1.0).contains(&self.mix) {
            return Err(format!("mix must be in [0.0, 1.0], got {}", self.mix));
        }
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Second-order IIR biquad (direct form II transposed)
// ---------------------------------------------------------------------------

/// Biquad filter state and coefficients.
#[derive(Debug, Clone)]
struct Biquad {
    b0: f32,
    b1: f32,
    b2: f32,
    a1: f32,
    a2: f32,
    s1: f32,
    s2: f32,
}

impl Biquad {
    fn process(&mut self, x: f32) -> f32 {
        let y = self.b0 * x + self.s1;
        self.s1 = self.b1 * x - self.a1 * y + self.s2;
        self.s2 = self.b2 * x - self.a2 * y;
        y
    }

    fn reset(&mut self) {
        self.s1 = 0.0;
        self.s2 = 0.0;
    }

    /// Butterworth 2nd-order lowpass at `freq_hz`.
    fn butterworth_lowpass(freq_hz: f32, sample_rate: f32) -> Self {
        let w0 = 2.0 * PI * freq_hz / sample_rate;
        let cos_w0 = w0.cos();
        let sin_w0 = w0.sin();
        let alpha = sin_w0 / std::f32::consts::SQRT_2; // Q=0.7071 (Butterworth)
        let a0 = 1.0 + alpha;
        Self {
            b0: ((1.0 - cos_w0) / 2.0) / a0,
            b1: (1.0 - cos_w0) / a0,
            b2: ((1.0 - cos_w0) / 2.0) / a0,
            a1: (-2.0 * cos_w0) / a0,
            a2: (1.0 - alpha) / a0,
            s1: 0.0,
            s2: 0.0,
        }
    }

    /// Butterworth 2nd-order highpass at `freq_hz`.
    fn butterworth_highpass(freq_hz: f32, sample_rate: f32) -> Self {
        let w0 = 2.0 * PI * freq_hz / sample_rate;
        let cos_w0 = w0.cos();
        let sin_w0 = w0.sin();
        let alpha = sin_w0 / std::f32::consts::SQRT_2;
        let a0 = 1.0 + alpha;
        Self {
            b0: ((1.0 + cos_w0) / 2.0) / a0,
            b1: (-(1.0 + cos_w0)) / a0,
            b2: ((1.0 + cos_w0) / 2.0) / a0,
            a1: (-2.0 * cos_w0) / a0,
            a2: (1.0 - alpha) / a0,
            s1: 0.0,
            s2: 0.0,
        }
    }
}

// ---------------------------------------------------------------------------
// BassEnhancer
// ---------------------------------------------------------------------------

/// Psychoacoustic bass enhancer.
pub struct BassEnhancer {
    config: BassEnhancerConfig,
    /// Lowpass filter to isolate bass band.
    lp: Biquad,
    /// Highpass filter to extract harmonic overtones above crossover.
    hp: Biquad,
    sample_rate: f32,
}

impl BassEnhancer {
    /// Create a new bass enhancer.
    ///
    /// # Panics
    ///
    /// Does not panic; frequency is clamped to a safe range internally.
    #[must_use]
    pub fn new(config: BassEnhancerConfig, sample_rate: f32) -> Self {
        let freq = config.frequency_hz.clamp(20.0, sample_rate * 0.45);
        let lp = Biquad::butterworth_lowpass(freq, sample_rate);
        let hp = Biquad::butterworth_highpass(freq, sample_rate);
        Self {
            config,
            lp,
            hp,
            sample_rate,
        }
    }

    /// Get a reference to the current configuration.
    #[must_use]
    pub fn config(&self) -> &BassEnhancerConfig {
        &self.config
    }

    /// Update the configuration and reinitialise filters.
    pub fn set_config(&mut self, config: BassEnhancerConfig) {
        let freq = config.frequency_hz.clamp(20.0, self.sample_rate * 0.45);
        self.lp = Biquad::butterworth_lowpass(freq, self.sample_rate);
        self.hp = Biquad::butterworth_highpass(freq, self.sample_rate);
        self.config = config;
    }

    /// Process a single sample.
    pub fn process_sample(&mut self, input: f32) -> f32 {
        // Extract bass band
        let bass = self.lp.process(input);

        // Waveshaper: generate harmonics
        let harmonic_content = self.generate_harmonics(bass);

        // Highpass the harmonic content to isolate overtones only
        let harmonics_hp = self.hp.process(harmonic_content);

        // Mix harmonics back with original signal
        input + harmonics_hp * self.config.mix
    }

    /// Process a buffer of samples in-place.
    pub fn process(&mut self, buffer: &mut [f32]) {
        for sample in buffer.iter_mut() {
            *sample = self.process_sample(*sample);
        }
    }

    /// Reset filter state.
    pub fn reset(&mut self) {
        self.lp.reset();
        self.hp.reset();
    }

    /// Generate harmonic content from the bass signal via waveshaping.
    fn generate_harmonics(&self, x: f32) -> f32 {
        if x.abs() < 1e-10 {
            return 0.0;
        }
        let drive = self.config.drive;
        let harmonics = self.config.harmonics;

        // Soft-clip the signal with progressive drive for each harmonic
        let mut out = 0.0f32;
        for n in 1..=harmonics {
            let gain = drive * (1.0 / n as f32);
            let driven = x * gain * (1 + n) as f32;
            // tanh waveshaper: produces odd+even harmonics at higher orders
            let shaped = if driven.abs() < 20.0 {
                driven.tanh()
            } else {
                driven.signum()
            };
            // Weight by inverse harmonic number (harmonic rolloff)
            out += shaped / n as f32;
        }
        out * drive
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    // ---- Config tests -------------------------------------------------------

    #[test]
    fn test_config_default() {
        let c = BassEnhancerConfig::default();
        assert!(c.frequency_hz > 0.0);
        assert!(c.harmonics >= 1);
        assert!(c.drive >= 0.0 && c.drive <= 1.0);
        assert!(c.mix >= 0.0 && c.mix <= 1.0);
        assert!(c.validate().is_ok());
    }

    #[test]
    fn test_config_gentle() {
        let c = BassEnhancerConfig::gentle();
        assert!(c.validate().is_ok());
    }

    #[test]
    fn test_config_aggressive() {
        let c = BassEnhancerConfig::aggressive();
        assert!(c.validate().is_ok());
    }

    #[test]
    fn test_config_validate_bad_frequency() {
        let c = BassEnhancerConfig {
            frequency_hz: 0.0,
            ..Default::default()
        };
        assert!(c.validate().is_err());
    }

    #[test]
    fn test_config_validate_zero_harmonics() {
        let c = BassEnhancerConfig {
            harmonics: 0,
            ..Default::default()
        };
        assert!(c.validate().is_err());
    }

    #[test]
    fn test_config_validate_bad_drive() {
        let c = BassEnhancerConfig {
            drive: 1.5,
            ..Default::default()
        };
        assert!(c.validate().is_err());
        let c2 = BassEnhancerConfig {
            drive: -0.1,
            ..Default::default()
        };
        assert!(c2.validate().is_err());
    }

    #[test]
    fn test_config_validate_bad_mix() {
        let c = BassEnhancerConfig {
            mix: 2.0,
            ..Default::default()
        };
        assert!(c.validate().is_err());
    }

    // ---- Construction -------------------------------------------------------

    #[test]
    fn test_new_does_not_panic() {
        let _ = BassEnhancer::new(BassEnhancerConfig::default(), 48_000.0);
    }

    #[test]
    fn test_new_with_extreme_frequency() {
        // Should clamp and not panic
        let c = BassEnhancerConfig {
            frequency_hz: 100_000.0,
            ..Default::default()
        };
        let _ = BassEnhancer::new(c, 44_100.0);
    }

    // ---- Processing --------------------------------------------------------

    #[test]
    fn test_process_sample_finite() {
        let mut e = BassEnhancer::new(BassEnhancerConfig::default(), 48_000.0);
        for i in 0..1024 {
            let s = (i as f32 * 0.01).sin() * 0.5;
            let out = e.process_sample(s);
            assert!(out.is_finite(), "output not finite at sample {i}");
        }
    }

    #[test]
    fn test_process_silent_input() {
        let mut e = BassEnhancer::new(BassEnhancerConfig::default(), 48_000.0);
        let out = e.process_sample(0.0);
        assert!((out).abs() < 1e-6);
    }

    #[test]
    fn test_process_buffer_length_preserved() {
        let mut e = BassEnhancer::new(BassEnhancerConfig::default(), 48_000.0);
        let mut buf = vec![0.1f32; 256];
        e.process(&mut buf);
        assert_eq!(buf.len(), 256);
    }

    #[test]
    fn test_process_buffer_all_finite() {
        let mut e = BassEnhancer::new(BassEnhancerConfig::default(), 48_000.0);
        let mut buf: Vec<f32> = (0..512)
            .map(|i| (i as f32 * 2.0 * PI / 512.0 * 80.0).sin() * 0.3)
            .collect();
        e.process(&mut buf);
        for &v in &buf {
            assert!(v.is_finite());
        }
    }

    #[test]
    fn test_zero_mix_passes_dry() {
        let config = BassEnhancerConfig {
            mix: 0.0,
            ..Default::default()
        };
        let mut e = BassEnhancer::new(config, 48_000.0);
        let input = 0.5f32;
        let out = e.process_sample(input);
        // With mix=0, no harmonics are added; lowpass filter state may cause
        // slight deviation at first sample but should be bounded.
        assert!(out.is_finite());
    }

    #[test]
    fn test_harmonics_add_energy() {
        // With high drive and mix, enhancer should produce output with energy
        // different from input (harmonics added)
        let config = BassEnhancerConfig {
            harmonics: 4,
            drive: 0.8,
            mix: 0.9,
            ..Default::default()
        };
        let mut e = BassEnhancer::new(config, 48_000.0);
        // Warm up the filter
        for _ in 0..100 {
            e.process_sample(0.3);
        }
        let out = e.process_sample(0.3f32);
        // Just ensure it's finite and non-zero
        assert!(out.is_finite());
        assert!(out.abs() > 0.0);
    }

    #[test]
    fn test_reset_clears_state() {
        let mut e = BassEnhancer::new(BassEnhancerConfig::default(), 48_000.0);
        for _ in 0..100 {
            e.process_sample(0.5);
        }
        e.reset();
        // After reset, first output should be close to input (no filter history)
        let out = e.process_sample(0.5);
        assert!(out.is_finite());
    }

    #[test]
    fn test_set_config_updates() {
        let mut e = BassEnhancer::new(BassEnhancerConfig::default(), 48_000.0);
        let new_cfg = BassEnhancerConfig {
            frequency_hz: 120.0,
            ..Default::default()
        };
        e.set_config(new_cfg.clone());
        assert!((e.config().frequency_hz - 120.0).abs() < 1e-5);
    }

    // ---- Biquad filter tests -----------------------------------------------

    #[test]
    fn test_biquad_lowpass_dc_passes() {
        let mut f = Biquad::butterworth_lowpass(100.0, 48_000.0);
        // Feed DC for many samples; output should stabilise close to input
        let mut out = 0.0f32;
        for _ in 0..1000 {
            out = f.process(1.0);
        }
        assert!((out - 1.0).abs() < 0.01, "DC not passing lowpass: {out}");
    }

    #[test]
    fn test_biquad_highpass_dc_blocked() {
        let mut f = Biquad::butterworth_highpass(100.0, 48_000.0);
        let mut out = 0.0f32;
        for _ in 0..1000 {
            out = f.process(1.0);
        }
        assert!(out.abs() < 0.01, "DC not blocked by highpass: {out}");
    }
}