audiobook-creation-exchange 0.1.0

//! Warmth shelving equaliser — two biquad IIR filters applied in a single pass.
//!
//! A low-frequency shelf boosts body around 180 Hz; a high-frequency shelf lifts
//! presence/air around 5 kHz. Both are gentle (≤ 3 dB) so the pipeline limiter
//! is not materially stressed.
//!
//! Biquad coefficients follow the Audio EQ Cookbook (Robert Bristow-Johnson).

/// Default low-shelf boost in dB.
pub const DEFAULT_LOW_SHELF_DB: f32 = 2.0;
/// Default low-shelf corner frequency in Hz.
pub const DEFAULT_LOW_SHELF_HZ: f32 = 180.0;
/// Default high-shelf boost in dB.
pub const DEFAULT_HIGH_SHELF_DB: f32 = 1.5;
/// Default high-shelf corner frequency in Hz.
pub const DEFAULT_HIGH_SHELF_HZ: f32 = 5_000.0;

/// Apply warmth EQ using default parameters.
pub fn apply_warmth(samples: &mut [i16], sample_rate: u32) {
    apply_warmth_with_params(
        samples,
        sample_rate,
        DEFAULT_LOW_SHELF_DB,
        DEFAULT_LOW_SHELF_HZ,
        DEFAULT_HIGH_SHELF_DB,
        DEFAULT_HIGH_SHELF_HZ,
    );
}

/// Apply warmth EQ with explicit shelf parameters.
pub fn apply_warmth_with_params(
    samples: &mut [i16],
    sample_rate: u32,
    low_shelf_db: f32,
    low_freq_hz: f32,
    high_shelf_db: f32,
    high_freq_hz: f32,
) {
    if samples.is_empty() || sample_rate == 0 {
        return;
    }

    let sr = sample_rate as f32;
    let (lo_b0, lo_b1, lo_b2, lo_a1, lo_a2) = low_shelf_coeffs(low_freq_hz, low_shelf_db, sr);
    let (hi_b0, hi_b1, hi_b2, hi_a1, hi_a2) = high_shelf_coeffs(high_freq_hz, high_shelf_db, sr);

    // Per-filter state (Direct Form I: x[n-1], x[n-2], y[n-1], y[n-2])
    let mut lo_x1 = 0f32;
    let mut lo_x2 = 0f32;
    let mut lo_y1 = 0f32;
    let mut lo_y2 = 0f32;

    let mut hi_x1 = 0f32;
    let mut hi_x2 = 0f32;
    let mut hi_y1 = 0f32;
    let mut hi_y2 = 0f32;

    for s in samples.iter_mut() {
        let x = *s as f32;

        // Low shelf
        let lo_y = lo_b0 * x + lo_b1 * lo_x1 + lo_b2 * lo_x2 - lo_a1 * lo_y1 - lo_a2 * lo_y2;
        lo_x2 = lo_x1;
        lo_x1 = x;
        lo_y2 = lo_y1;
        lo_y1 = lo_y;

        // High shelf
        let hi_y = hi_b0 * lo_y + hi_b1 * hi_x1 + hi_b2 * hi_x2 - hi_a1 * hi_y1 - hi_a2 * hi_y2;
        hi_x2 = hi_x1;
        hi_x1 = lo_y;
        hi_y2 = hi_y1;
        hi_y1 = hi_y;

        *s = hi_y.round().clamp(i16::MIN as f32, i16::MAX as f32) as i16;
    }
}

/// Low-shelf biquad coefficients.
/// Returns (b0, b1, b2, a1, a2) — normalised so a0 = 1.
fn low_shelf_coeffs(freq_hz: f32, gain_db: f32, sr: f32) -> (f32, f32, f32, f32, f32) {
    let a = 10f32.powf(gain_db / 40.0); // sqrt(A)
    let w0 = 2.0 * std::f32::consts::PI * freq_hz / sr;
    let cos_w0 = w0.cos();
    let sin_w0 = w0.sin();
    let alpha = sin_w0 / 2.0 * (a + 1.0 / a).sqrt(); // slope = 1

    let a0 = (a + 1.0) + (a - 1.0) * cos_w0 + 2.0 * a.sqrt() * alpha;
    let b0 = a * ((a + 1.0) - (a - 1.0) * cos_w0 + 2.0 * a.sqrt() * alpha) / a0;
    let b1 = 2.0 * a * ((a - 1.0) - (a + 1.0) * cos_w0) / a0;
    let b2 = a * ((a + 1.0) - (a - 1.0) * cos_w0 - 2.0 * a.sqrt() * alpha) / a0;
    let a1 = -2.0 * ((a - 1.0) + (a + 1.0) * cos_w0) / a0;
    let a2 = ((a + 1.0) + (a - 1.0) * cos_w0 - 2.0 * a.sqrt() * alpha) / a0;

    (b0, b1, b2, a1, a2)
}

/// High-shelf biquad coefficients.
/// Returns (b0, b1, b2, a1, a2) — normalised so a0 = 1.
fn high_shelf_coeffs(freq_hz: f32, gain_db: f32, sr: f32) -> (f32, f32, f32, f32, f32) {
    let a = 10f32.powf(gain_db / 40.0);
    let w0 = 2.0 * std::f32::consts::PI * freq_hz / sr;
    let cos_w0 = w0.cos();
    let sin_w0 = w0.sin();
    let alpha = sin_w0 / 2.0 * (a + 1.0 / a).sqrt();

    let a0 = (a + 1.0) - (a - 1.0) * cos_w0 + 2.0 * a.sqrt() * alpha;
    let b0 = a * ((a + 1.0) + (a - 1.0) * cos_w0 + 2.0 * a.sqrt() * alpha) / a0;
    let b1 = -2.0 * a * ((a - 1.0) + (a + 1.0) * cos_w0) / a0;
    let b2 = a * ((a + 1.0) + (a - 1.0) * cos_w0 - 2.0 * a.sqrt() * alpha) / a0;
    let a1 = 2.0 * ((a - 1.0) - (a + 1.0) * cos_w0) / a0;
    let a2 = ((a + 1.0) - (a - 1.0) * cos_w0 - 2.0 * a.sqrt() * alpha) / a0;

    (b0, b1, b2, a1, a2)
}

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

    const SR: u32 = 24_000;

    fn pure_tone(freq_hz: f32, amplitude: f32, secs: f32) -> Vec<i16> {
        let n = (SR as f32 * secs) as usize;
        (0..n)
            .map(|i| {
                let v =
                    amplitude * (2.0 * std::f32::consts::PI * freq_hz * i as f32 / SR as f32).sin();
                v.clamp(i16::MIN as f32, i16::MAX as f32) as i16
            })
            .collect()
    }

    #[test]
    fn empty_input_is_a_no_op() {
        let mut samples: Vec<i16> = Vec::new();
        apply_warmth(&mut samples, SR); // must not panic
    }

    #[test]
    fn low_shelf_boosts_bass_tone() {
        let original = pure_tone(100.0, 6_000.0, 0.5);
        let mut boosted = original.clone();
        apply_warmth_with_params(&mut boosted, SR, 3.0, 180.0, 0.0, 5_000.0);

        let orig_rms: f32 = rms_db(&original);
        let proc_rms: f32 = rms_db(&boosted);
        assert!(
            proc_rms > orig_rms,
            "Low shelf should boost bass tone: before={:.2} after={:.2}",
            orig_rms,
            proc_rms
        );
    }

    #[test]
    fn high_shelf_boosts_presence_tone() {
        let original = pure_tone(8_000.0, 6_000.0, 0.5);
        let mut boosted = original.clone();
        apply_warmth_with_params(&mut boosted, SR, 0.0, 180.0, 3.0, 5_000.0);

        let orig_rms: f32 = rms_db(&original);
        let proc_rms: f32 = rms_db(&boosted);
        assert!(
            proc_rms > orig_rms,
            "High shelf should boost presence tone: before={:.2} after={:.2}",
            orig_rms,
            proc_rms
        );
    }

    #[test]
    fn midband_tone_is_unaffected() {
        // 1 kHz is between the two shelves — should pass through with < 0.5 dB change.
        let original = pure_tone(1_000.0, 8_000.0, 0.5);
        let mut processed = original.clone();
        apply_warmth(&mut processed, SR);

        let orig_rms = rms_db(&original);
        let proc_rms = rms_db(&processed);
        let diff_db = proc_rms - orig_rms;
        assert!(diff_db.abs() < 0.5, "Mid-band altered by {:.2} dB", diff_db);
    }

    #[test]
    fn boost_does_not_clip_headroom() {
        // Use amplitude 20 000 (below full scale) — shelves add ≤ 3 dB, still safe.
        let mut samples = pure_tone(100.0, 20_000.0, 0.3);
        apply_warmth(&mut samples, SR);
        let max = samples.iter().map(|&s| s.abs()).max().unwrap_or(0);
        assert!(max <= i16::MAX, "Warmth EQ clipped: max sample = {}", max);
    }
}