zk-audio 0.1.0

Audio processing library for voice recording and enhancement
Documentation
use super::support::smooth_towards;
use super::{analyze_frame, classify_frame, FrameClass};
use crate::contracts::AudioProcessor;
use crate::core::{AudioError, AudioFrame, AudioProfile, AudioResult, AudioSpec};
use crate::profiles::profile_tuning;
use realfft::{num_complex::Complex32, ComplexToReal, RealFftPlanner, RealToComplex};
use std::sync::Arc;

struct ResonancePlan {
    len: usize,
    hop_len: usize,
    forward: Arc<dyn RealToComplex<f32>>,
    inverse: Arc<dyn ComplexToReal<f32>>,
    analysis_window: Vec<f32>,
    synthesis_window: Vec<f32>,
    input: Vec<f32>,
    spectrum: Vec<Complex32>,
    output: Vec<f32>,
    prev_input_tail: Vec<f32>,
    overlap_output: Vec<f32>,
}

impl ResonancePlan {
    fn new(hop_len: usize) -> Self {
        let len = (hop_len.max(64) * 2).next_power_of_two();
        let hop_len = len / 2;
        let mut planner = RealFftPlanner::<f32>::new();
        let forward = planner.plan_fft_forward(len);
        let inverse = planner.plan_fft_inverse(len);
        let analysis_window = (0..len)
            .map(|i| {
                let position = i as f32 / (len.saturating_sub(1).max(1)) as f32;
                (0.5 - 0.5 * (std::f32::consts::TAU * position).cos()).sqrt()
            })
            .collect::<Vec<_>>();
        let synthesis_window = analysis_window.clone();

        let input = forward.make_input_vec();
        let spectrum = forward.make_output_vec();
        let output = inverse.make_output_vec();
        let prev_input_tail = vec![0.0; hop_len];
        let overlap_output = vec![0.0; hop_len];

        Self {
            len,
            hop_len,
            forward,
            inverse,
            analysis_window,
            synthesis_window,
            input,
            spectrum,
            output,
            prev_input_tail,
            overlap_output,
        }
    }
}

pub struct MidResonanceSuppressorProcessor {
    profile: AudioProfile,
    sample_rate: u32,
    noise_floor: f32,
    current_strength: f32,
    plan: Option<ResonancePlan>,
}

impl MidResonanceSuppressorProcessor {
    pub fn new(profile: AudioProfile) -> Self {
        Self {
            profile,
            sample_rate: 44_100,
            noise_floor: 0.01,
            current_strength: 0.0,
            plan: None,
        }
    }

    fn ensure_plan(&mut self, hop_len: usize) {
        let needs_rebuild = self
            .plan
            .as_ref()
            .map(|plan| plan.hop_len != hop_len)
            .unwrap_or(true);
        if needs_rebuild {
            self.plan = Some(ResonancePlan::new(hop_len));
        }
    }
}

impl AudioProcessor for MidResonanceSuppressorProcessor {
    fn name(&self) -> &'static str {
        "mid_resonance_suppressor"
    }

    fn prepare(&mut self, spec: AudioSpec) -> AudioResult<()> {
        self.sample_rate = spec.sample_rate;
        self.noise_floor = 0.01;
        self.current_strength = 0.0;
        Ok(())
    }

    fn process(&mut self, frame: &mut AudioFrame) -> AudioResult<()> {
        if frame.samples.len() < 32 {
            return Ok(());
        }

        self.ensure_plan(frame.samples.len());
        let plan = self
            .plan
            .as_mut()
            .ok_or_else(|| AudioError::new("Mid resonance plan missing"))?;
        if frame.samples.len() != plan.hop_len {
            return Ok(());
        }

        let features = analyze_frame(&frame.samples, self.noise_floor);
        let frame_class = classify_frame(self.profile, features, self.noise_floor);
        let rms = features.rms;
        let peak = features.peak;
        if peak < self.noise_floor * 3.0 && rms < self.noise_floor * 2.0 {
            self.noise_floor = self.noise_floor * 0.94 + rms.max(0.0005) * 0.06;
        } else {
            self.noise_floor = self.noise_floor * 0.997 + rms.max(0.0005) * 0.003;
        }

        let tuning = profile_tuning(self.profile);
        let base_strength = match frame_class {
            FrameClass::NoiseOnly => tuning.mid_resonance_strength,
            FrameClass::Transitional => tuning.mid_resonance_strength * 0.92,
            FrameClass::SpeechLike => tuning.mid_resonance_strength * 0.42,
        };
        let speech_relief = ((features.speechiness - tuning.adaptive_min_speechiness)
            / tuning.adaptive_min_speechiness.max(1.0))
        .clamp(0.0, 1.0);
        let target_strength = (base_strength * (1.0 - 0.28 * speech_relief)).clamp(0.0, 0.9);
        let smoothing = if target_strength > self.current_strength {
            tuning.mid_resonance_attack
        } else {
            tuning.mid_resonance_release
        };
        self.current_strength = smooth_towards(self.current_strength, target_strength, smoothing);

        let current_input = frame.samples.clone();
        for (index, value) in plan.input.iter_mut().enumerate() {
            let sample = if index < plan.hop_len {
                plan.prev_input_tail[index]
            } else {
                current_input[index - plan.hop_len]
            };
            *value = sample * plan.analysis_window[index];
        }

        plan.forward
            .process(&mut plan.input, &mut plan.spectrum)
            .map_err(|err| AudioError::new(format!("Mid resonance forward FFT failed: {}", err)))?;

        let magnitudes = plan
            .spectrum
            .iter()
            .map(|bin| bin.norm())
            .collect::<Vec<_>>();
        let bin_hz = self.sample_rate as f32 / plan.len as f32;
        for index in 3..plan.spectrum.len().saturating_sub(3) {
            let frequency_hz = index as f32 * bin_hz;
            if frequency_hz < tuning.mid_resonance_low_hz
                || frequency_hz > tuning.mid_resonance_high_hz
            {
                continue;
            }

            let magnitude = magnitudes[index];
            let left_avg =
                (magnitudes[index - 1] + magnitudes[index - 2] + magnitudes[index - 3]) / 3.0;
            let right_avg =
                (magnitudes[index + 1] + magnitudes[index + 2] + magnitudes[index + 3]) / 3.0;
            let neighborhood = ((left_avg + right_avg) * 0.5).max(1e-5);
            let prominence = magnitude / neighborhood;
            if prominence <= tuning.mid_resonance_prominence {
                continue;
            }

            let narrowness = ((prominence - tuning.mid_resonance_prominence)
                / tuning.mid_resonance_prominence.max(0.2))
            .clamp(0.0, 1.0);
            let attenuation = (self.current_strength * narrowness).clamp(0.0, 0.82);
            plan.spectrum[index] *= 1.0 - attenuation;
            if index > 0 {
                plan.spectrum[index - 1] *= 1.0 - attenuation * 0.24;
            }
            if index + 1 < plan.spectrum.len() {
                plan.spectrum[index + 1] *= 1.0 - attenuation * 0.24;
            }
        }

        plan.inverse
            .process(&mut plan.spectrum, &mut plan.output)
            .map_err(|err| AudioError::new(format!("Mid resonance inverse FFT failed: {}", err)))?;

        let normalize = 1.0 / plan.len as f32;
        for index in 0..plan.hop_len {
            let current = plan.output[index] * normalize * plan.synthesis_window[index];
            frame.samples[index] = current + plan.overlap_output[index];
        }
        for index in 0..plan.hop_len {
            plan.overlap_output[index] = plan.output[index + plan.hop_len]
                * normalize
                * plan.synthesis_window[index + plan.hop_len];
            plan.prev_input_tail[index] = current_input[index];
        }

        Ok(())
    }
}

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

    fn test_spec() -> AudioSpec {
        AudioSpec {
            sample_rate: 44_100,
            channels: 1,
        }
    }

    #[test]
    fn mid_resonance_suppressor_reduces_narrow_mid_peak() {
        let spec = test_spec();
        let mut processor = MidResonanceSuppressorProcessor::new(AudioProfile::VoiceHvac);
        processor.prepare(spec).unwrap();

        let source = (0..512)
            .map(|i| {
                let t = i as f32 / spec.sample_rate as f32;
                let body = (std::f32::consts::TAU * 220.0 * t).sin() * 0.10;
                let ring = (std::f32::consts::TAU * 1_850.0 * t).sin() * 0.06;
                body + ring
            })
            .collect::<Vec<_>>();
        let mut processed = source.clone();
        for _ in 0..6 {
            let mut frame = AudioFrame {
                samples: source.clone(),
                spec,
            };
            processor.process(&mut frame).unwrap();
            processed = frame.samples;
        }

        let mean_delta = processed
            .iter()
            .zip(source.iter())
            .map(|(after, before)| (after - before).abs())
            .sum::<f32>()
            / source.len().max(1) as f32;

        assert!(processor.current_strength > 0.01);
        assert!(mean_delta > 0.001);
    }

    #[test]
    fn mid_resonance_suppressor_is_light_on_voice_body() {
        let spec = test_spec();
        let mut processor = MidResonanceSuppressorProcessor::new(AudioProfile::VoiceClean);
        processor.prepare(spec).unwrap();

        let mut frame = AudioFrame {
            samples: (0..512)
                .map(|i| {
                    let t = i as f32 / spec.sample_rate as f32;
                    (std::f32::consts::TAU * 180.0 * t).sin() * 0.12
                })
                .collect(),
            spec,
        };

        let source = frame.samples.clone();
        processor.process(&mut frame).unwrap();
        let mean_delta = frame
            .samples
            .iter()
            .zip(source.iter())
            .map(|(after, before)| (after - before).abs())
            .sum::<f32>()
            / source.len() as f32;

        assert!(processor.current_strength < 0.12);
        assert!(mean_delta < 0.08);
    }
}