nsm/crates/player/src/spectrum.rs

use std::sync::Arc;

use realfft::num_complex::Complex;
use realfft::{RealFftPlanner, RealToComplex};

/// Размер окна FFT — степень двойки, разумный компромисс между частотным и временным
/// разрешением для визуализации (не для точного научного анализа).
pub const FFT_SIZE: usize = 1024;
/// Количество полос на выходе (логарифмическая шкала, как у обычных спектроанализаторов).
pub const SPECTRUM_BANDS: usize = 20;

pub struct SpectrumAnalyzer {
    fft: Arc<dyn RealToComplex<f32>>,
    window: Vec<f32>,
    input_scratch: Vec<f32>,
    output_scratch: Vec<Complex<f32>>,
}

impl SpectrumAnalyzer {
    pub fn new() -> Self {
        let mut planner = RealFftPlanner::<f32>::new();
        let fft = planner.plan_fft_forward(FFT_SIZE);
        // окно Ханна — сглаживает края окна, уменьшает "утечку" спектра между полосами
        let window: Vec<f32> = (0..FFT_SIZE)
            .map(|i| {
                0.5 - 0.5
                    * (2.0 * std::f32::consts::PI * i as f32 / (FFT_SIZE as f32 - 1.0)).cos()
            })
            .collect();
        let input_scratch = fft.make_input_vec();
        let output_scratch = fft.make_output_vec();
        SpectrumAnalyzer {
            fft,
            window,
            input_scratch,
            output_scratch,
        }
    }

    /// `samples` — ровно `FFT_SIZE` моно-сэмплов (последнее окно). Возвращает `SPECTRUM_BANDS`
    /// значений, сгруппированных по логарифмической шкале частот (как воспринимает слух),
    /// нормализованных примерно в диапазон 0..1.
    pub fn compute(&mut self, samples: &[f32], sample_rate: f32) -> [f32; SPECTRUM_BANDS] {
        debug_assert_eq!(samples.len(), FFT_SIZE);
        for i in 0..FFT_SIZE {
            self.input_scratch[i] = samples[i] * self.window[i];
        }
        if self
            .fft
            .process(&mut self.input_scratch, &mut self.output_scratch)
            .is_err()
        {
            return [0.0; SPECTRUM_BANDS];
        }

        let magnitudes: Vec<f32> = self
            .output_scratch
            .iter()
            .map(|c| (c.re * c.re + c.im * c.im).sqrt())
            .collect();

        bucket_log_scale(&magnitudes, sample_rate)
    }
}

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

/// Группирует линейные бины FFT в `SPECTRUM_BANDS` логарифмических полос от ~20Гц до
/// половины частоты дискретизации (или 20кГц, что меньше), и нормализует громким-сжатием
/// (sqrt) для более "живой" визуальной динамики на тихих сигналах.
fn bucket_log_scale(magnitudes: &[f32], sample_rate: f32) -> [f32; SPECTRUM_BANDS] {
    let min_freq: f32 = 20.0;
    let max_freq: f32 = (sample_rate / 2.0).min(20_000.0);
    let log_min = min_freq.ln();
    let log_max = max_freq.ln();
    let bin_hz = sample_rate / FFT_SIZE as f32;

    let mut bands = [0.0f32; SPECTRUM_BANDS];
    for (b, band) in bands.iter_mut().enumerate() {
        let f_lo = (log_min + (log_max - log_min) * b as f32 / SPECTRUM_BANDS as f32).exp();
        let f_hi = (log_min + (log_max - log_min) * (b + 1) as f32 / SPECTRUM_BANDS as f32).exp();
        let bin_lo = ((f_lo / bin_hz).round() as usize).max(1);
        let bin_hi = ((f_hi / bin_hz).round() as usize)
            .min(magnitudes.len().saturating_sub(1))
            .max(bin_lo + 1);

        let slice = &magnitudes[bin_lo.min(magnitudes.len())..bin_hi.min(magnitudes.len())];
        let avg = if slice.is_empty() {
            0.0
        } else {
            slice.iter().sum::<f32>() / slice.len() as f32
        };
        // нормализация: сырые магнитуды FFT могут быть большими числами без верхней границы
        // конкретного смысла для отображения — берём корень и делим на эмпирическую константу,
        // чтобб типичная громкая музыка укладывалась примерно в 0..1
        *band = (avg.sqrt() / 12.0).clamp(0.0, 1.0);
    }
    bands
}

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

    #[test]
    fn sine_wave_peaks_in_expected_band() {
        let sample_rate = 44100.0;
        // тестовый тон 1000 Гц — должен дать пик в средних полосах, а не в самых низких/высоких
        let freq = 1000.0f32;
        let samples: Vec<f32> = (0..FFT_SIZE)
            .map(|i| (2.0 * std::f32::consts::PI * freq * i as f32 / sample_rate).sin())
            .collect();

        let mut analyzer = SpectrumAnalyzer::new();
        let bands = analyzer.compute(&samples, sample_rate);

        let (peak_idx, peak_val) = bands
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
            .unwrap();

        // не самая первая (бас) и не самая последняя (высокие) полоса
        assert!(peak_idx > 2 && peak_idx < SPECTRUM_BANDS - 2, "peak at {peak_idx}");
        assert!(*peak_val > 0.05, "peak too quiet: {peak_val}");
    }

    #[test]
    fn silence_gives_near_zero_bands() {
        let samples = vec![0.0f32; FFT_SIZE];
        let mut analyzer = SpectrumAnalyzer::new();
        let bands = analyzer.compute(&samples, 44100.0);
        for b in bands {
            assert!(b < 0.01, "expected near-zero for silence, got {b}");
        }
    }
}