onedrop_engine/

audio.rs

//! Audio processing and analysis.
//!
//! The MilkDrop preset surface (per-frame / per-pixel equations, beat-detect
//! thresholds, custom waves' volume modulation) reads three split-band volumes:
//! `bass`, `mid`, `treb`. They must reflect *frequency content* — kick drums
//! push `bass`, cymbals push `treb` — not just overall loudness. The original
//! placeholder in this file did RMS over three time-domain chunks of the input
//! window, so all three bands tracked overall loudness and presets couldn't
//! distinguish low- from high-frequency energy. This analyzer now runs an FFT
//! and integrates the magnitude spectrum over the MD2-standard frequency
//! ranges (`FFTAnalyzer::get_bass` / `_mid` / `_treble`).

use crate::fft::FFTAnalyzer;
use onedrop_renderer::AudioLevels;

/// Maps the FFT-band averages to roughly the [0, 2-ish] range MD2 presets
/// expect (`bass > 1.5` is the conventional "loud kick" threshold for
/// `BeatDetectionMode::HardCut1`). A unit-amplitude sine in the bass range
/// produces an averaged magnitude of `~0.5 / num_bass_bins` (the FFT spreads
/// one bin's peak across the averaging window); with the default 1024-point
/// FFT @ 44.1 kHz that's ~0.08, so gain 12.5 maps a pure bass sine to ~1.0.
/// Real-world music with a bass kick lands in the ~0.5-2.0 range with `50.0`,
/// which lines up reasonably with MD2's expected dynamics. Tunable per-user
/// later via the GUI if needed.
const FFT_BAND_GAIN: f32 = 50.0;

/// Target band level the AGC steers each band toward. MD2 presets are
/// authored against the convention that `bass`/`mid`/`treb` average
/// ~1.0 on typical content; `bass > 1.5` then unambiguously means
/// "louder than usual" (kick / snare / hat) rather than just "loud
/// track". With AGC off, quiet tracks float around 0.2-0.4 and never
/// trigger beat-reactive presets; loud tracks pin at 8.0 (the clamp)
/// and beat detection saturates.
const AGC_TARGET_LEVEL: f32 = 1.0;

/// Smoothing factor for the AGC's per-band running average. Closer to
/// 1.0 = longer history. ~50-100 frames at 60 FPS feels right — slow
/// enough that a single kick doesn't dump the gain, fast enough that a
/// track change resettles within a couple of seconds.
const AGC_AVG_ALPHA: f32 = 0.99;

/// Clamp the AGC gain to a sane range. Both bounds matter: a track
/// with absolute silence on a band shouldn't let the gain run to
/// infinity (the floor at 0.5 caps the boost on near-silent bands),
/// and a band that's pure-tone-driven shouldn't crush below 1/10 of
/// the raw signal.
const AGC_GAIN_MIN: f32 = 0.1;
const AGC_GAIN_MAX: f32 = 10.0;

/// FFT window length. Powers of 2; 1024 gives ~43 Hz bin width @ 44.1 kHz,
/// which is enough granularity for the bass band (smaller windows under-resolve
/// 20-250 Hz; larger windows raise latency without buying much for the
/// 3-band split MD2 uses).
const FFT_SIZE: usize = 1024;

/// Audio analyzer for extracting per-band volumes.
pub struct AudioAnalyzer {
    fft: FFTAnalyzer,

    /// Long-term smoothed bass — the `bass_att` variable preset equations read.
    bass_att: f32,

    /// Long-term smoothed mid — the `mid_att` variable preset equations read.
    mid_att: f32,

    /// Long-term smoothed treble — the `treb_att` variable preset equations read.
    treb_att: f32,

    /// Exponential-smoothing retention factor for the `_att` channels. Closer
    /// to 1.0 = longer history. MD2's `_att` channels are heavily smoothed
    /// so they represent recent average loudness, not instantaneous.
    attenuation: f32,

    /// Per-band running average used by the AGC to steer the gain.
    /// Updated each frame from the raw pre-AGC band level; the AGC
    /// divides the target level by this to compute the gain factor.
    /// Floor at 0.05 so the gain doesn't blow up on near-silent input.
    agc_bass_avg: f32,
    agc_mid_avg: f32,
    agc_treb_avg: f32,
}

impl AudioAnalyzer {
    /// Create a new audio analyzer.
    pub fn new(sample_rate: f32) -> Self {
        Self {
            fft: FFTAnalyzer::new_or_default(FFT_SIZE, sample_rate),
            bass_att: 0.0,
            mid_att: 0.0,
            treb_att: 0.0,
            attenuation: 0.93,
            // Seed AGC averages at the target so the gain starts at 1.0
            // (no boost on the first frame); converges to the actual
            // long-term level over the next ~100 frames.
            agc_bass_avg: AGC_TARGET_LEVEL,
            agc_mid_avg: AGC_TARGET_LEVEL,
            agc_treb_avg: AGC_TARGET_LEVEL,
        }
    }

    /// Analyze audio samples and return per-band volumes.
    ///
    /// Pipeline: FFT → raw band integral → AGC. The AGC keeps each
    /// band hovering near `AGC_TARGET_LEVEL = 1.0` on typical content,
    /// so `bass > 1.5` on the output unambiguously means "louder than
    /// usual" — the convention every MD2 preset's beat-detection block
    /// is authored against. Without the AGC, quiet tracks float around
    /// 0.2-0.4 and never trigger `BeatDetectionMode::HardCut1`; loud
    /// tracks pin at the 8.0 clamp and beat detection saturates.
    pub fn analyze(&mut self, samples: &[f32]) -> AudioLevels {
        self.fft.analyze(samples);

        let raw_bass = (self.fft.get_bass() * FFT_BAND_GAIN).clamp(0.0, 8.0);
        let raw_mid = (self.fft.get_mid() * FFT_BAND_GAIN).clamp(0.0, 8.0);
        let raw_treb = (self.fft.get_treble() * FFT_BAND_GAIN).clamp(0.0, 8.0);

        // Update the AGC running averages and derive per-band gains.
        // Floor at 0.05 so the divide doesn't explode on silent bands.
        self.agc_bass_avg = self.agc_bass_avg * AGC_AVG_ALPHA + raw_bass * (1.0 - AGC_AVG_ALPHA);
        self.agc_mid_avg = self.agc_mid_avg * AGC_AVG_ALPHA + raw_mid * (1.0 - AGC_AVG_ALPHA);
        self.agc_treb_avg = self.agc_treb_avg * AGC_AVG_ALPHA + raw_treb * (1.0 - AGC_AVG_ALPHA);

        let bass_gain =
            (AGC_TARGET_LEVEL / self.agc_bass_avg.max(0.05)).clamp(AGC_GAIN_MIN, AGC_GAIN_MAX);
        let mid_gain =
            (AGC_TARGET_LEVEL / self.agc_mid_avg.max(0.05)).clamp(AGC_GAIN_MIN, AGC_GAIN_MAX);
        let treb_gain =
            (AGC_TARGET_LEVEL / self.agc_treb_avg.max(0.05)).clamp(AGC_GAIN_MIN, AGC_GAIN_MAX);

        let bass = (raw_bass * bass_gain).clamp(0.0, 8.0);
        let mid = (raw_mid * mid_gain).clamp(0.0, 8.0);
        let treb = (raw_treb * treb_gain).clamp(0.0, 8.0);

        self.bass_att = self.bass_att * self.attenuation + bass * (1.0 - self.attenuation);
        self.mid_att = self.mid_att * self.attenuation + mid * (1.0 - self.attenuation);
        self.treb_att = self.treb_att * self.attenuation + treb * (1.0 - self.attenuation);

        AudioLevels {
            bass,
            mid,
            treb,
            bass_att: self.bass_att,
            mid_att: self.mid_att,
            treb_att: self.treb_att,
        }
    }

    /// Set attenuation factor.
    pub fn set_attenuation(&mut self, attenuation: f32) {
        self.attenuation = attenuation.clamp(0.0, 1.0);
    }

    /// Reset attenuated values.
    pub fn reset(&mut self) {
        self.bass_att = 0.0;
        self.mid_att = 0.0;
        self.treb_att = 0.0;
        // Re-seed AGC averages at the target so the first frame after
        // reset doesn't start at full boost.
        self.agc_bass_avg = AGC_TARGET_LEVEL;
        self.agc_mid_avg = AGC_TARGET_LEVEL;
        self.agc_treb_avg = AGC_TARGET_LEVEL;
    }
}

impl Default for AudioAnalyzer {
    fn default() -> Self {
        Self::new(44100.0)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;
    use std::f32::consts::PI;

    /// Build a unit-amplitude sine at the given Hz, sampled at 44.1 kHz.
    fn sine(freq_hz: f32, n: usize) -> Vec<f32> {
        (0..n)
            .map(|i| (2.0 * PI * freq_hz * i as f32 / 44100.0).sin())
            .collect()
    }

    #[test]
    fn test_analyze_silence() {
        let mut analyzer = AudioAnalyzer::new(44100.0);
        let samples = vec![0.0; 1024];

        let levels = analyzer.analyze(&samples);

        assert_relative_eq!(levels.bass, 0.0, epsilon = 0.01);
        assert_relative_eq!(levels.mid, 0.0, epsilon = 0.01);
        assert_relative_eq!(levels.treb, 0.0, epsilon = 0.01);
    }

    /// Frequency separation matters: a 60 Hz tone should drive `bass`
    /// dominantly, a 1 kHz tone `mid`, an 8 kHz tone `treb`. This is the
    /// regression test for the old time-chunked-RMS placeholder, which
    /// reported all three bands as roughly equal for any input.
    #[test]
    fn bass_tone_drives_bass_band_dominantly() {
        let mut analyzer = AudioAnalyzer::new(44100.0);
        let levels = analyzer.analyze(&sine(60.0, 1024));
        assert!(levels.bass > 5.0 * levels.mid, "{:?}", levels);
        assert!(levels.bass > 5.0 * levels.treb, "{:?}", levels);
    }

    #[test]
    fn mid_tone_drives_mid_band_dominantly() {
        let mut analyzer = AudioAnalyzer::new(44100.0);
        let levels = analyzer.analyze(&sine(1000.0, 1024));
        assert!(levels.mid > 5.0 * levels.bass, "{:?}", levels);
        assert!(levels.mid > 5.0 * levels.treb, "{:?}", levels);
    }

    #[test]
    fn treble_tone_drives_treb_band_dominantly() {
        let mut analyzer = AudioAnalyzer::new(44100.0);
        let levels = analyzer.analyze(&sine(8000.0, 1024));
        assert!(levels.treb > 5.0 * levels.bass, "{:?}", levels);
        assert!(levels.treb > 5.0 * levels.mid, "{:?}", levels);
    }

    /// A unit-amplitude bass sine is the loudest signal cpal will hand us.
    /// Real music sits ~10-20 dB below this, so a unit sine producing
    /// ~5 means typical music produces bass in the [0.5, 2.0] range MD2
    /// presets expect for "moderate" → "loud kick". Pins the gain.
    #[test]
    fn unit_bass_sine_is_within_clamp_and_above_floor() {
        let mut analyzer = AudioAnalyzer::new(44100.0);
        let levels = analyzer.analyze(&sine(80.0, 1024));
        assert!(levels.bass > 1.0, "bass too quiet: {:?}", levels.bass);
        assert!(levels.bass < 8.0, "bass clamp violated: {:?}", levels.bass);
    }

    /// The AGC must steer the per-band level toward `AGC_TARGET_LEVEL`
    /// on sustained input: a quiet but steady bass tone should boost
    /// over many frames until the reported `bass` lands near 1.0,
    /// independent of the source amplitude. Without this, presets'
    /// `bass > 1.5` beat threshold never fires on quiet tracks.
    #[test]
    fn agc_steers_quiet_band_toward_target() {
        let mut analyzer = AudioAnalyzer::new(44100.0);
        let quiet: Vec<f32> = sine(60.0, 1024).into_iter().map(|s| s * 0.1).collect();

        // First frame: low signal × seed-gain 1.0 = low reported bass.
        let first = analyzer.analyze(&quiet);
        assert!(
            first.bass < 1.0,
            "first frame should not be boosted: {:?}",
            first.bass
        );

        // After many frames, the AGC running average converges to the
        // quiet level and the gain inverts it back to ~target.
        for _ in 0..1000 {
            analyzer.analyze(&quiet);
        }
        let settled = analyzer.analyze(&quiet);
        assert!(
            settled.bass > 0.7,
            "AGC should boost quiet sustained tone toward target, got {}",
            settled.bass
        );
    }

    /// Inverse of the boost case: a LOUD sustained tone must be
    /// attenuated by the AGC so the reported level stays near the
    /// target. Otherwise loud tracks pin the clamp at 8.0 and beat
    /// detection saturates.
    #[test]
    fn agc_attenuates_loud_band_toward_target() {
        let mut analyzer = AudioAnalyzer::new(44100.0);
        let loud: Vec<f32> = sine(60.0, 1024).into_iter().map(|s| s * 4.0).collect();

        for _ in 0..1000 {
            analyzer.analyze(&loud);
        }
        let settled = analyzer.analyze(&loud);
        assert!(
            settled.bass < 4.0,
            "AGC should attenuate loud sustained tone, got {}",
            settled.bass
        );
    }

    #[test]
    fn att_channels_smooth_across_frames() {
        let mut analyzer = AudioAnalyzer::new(44100.0);
        let samples = sine(1000.0, 1024);

        let first = analyzer.analyze(&samples);
        // After one frame, _att has barely accumulated (long smoothing).
        assert!(first.mid_att < first.mid);

        // After many frames, _att approaches the instantaneous value.
        for _ in 0..200 {
            analyzer.analyze(&samples);
        }
        let settled = analyzer.analyze(&samples);
        assert_relative_eq!(settled.mid_att, settled.mid, epsilon = 0.05);
    }
}