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 TailPlan {
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>,
tail_envelope: Vec<f32>,
}
impl TailPlan {
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];
let tail_envelope = vec![0.0; spectrum.len()];
Self {
len,
hop_len,
forward,
inverse,
analysis_window,
synthesis_window,
input,
spectrum,
output,
prev_input_tail,
overlap_output,
tail_envelope,
}
}
}
pub struct TailSuppressionProcessor {
profile: AudioProfile,
sample_rate: u32,
noise_floor: f32,
current_strength: f32,
plan: Option<TailPlan>,
}
impl TailSuppressionProcessor {
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(TailPlan::new(hop_len));
}
}
}
impl AudioProcessor for TailSuppressionProcessor {
fn name(&self) -> &'static str {
"tail_suppression"
}
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("Tail suppression 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.998 + rms.max(0.0005) * 0.002;
}
let tuning = profile_tuning(self.profile);
let base_strength = match frame_class {
FrameClass::NoiseOnly => tuning.tail_strength * 0.15,
FrameClass::Transitional => tuning.tail_strength * 0.85,
FrameClass::SpeechLike => tuning.tail_strength,
}
.clamp(0.0, 0.95);
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.25 * speech_relief)).clamp(0.0, 0.95);
let smoothing = if target_strength > self.current_strength {
tuning.tail_attack
} else {
tuning.tail_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!("Tail suppression forward FFT failed: {}", err))
})?;
let bin_hz = self.sample_rate as f32 / plan.len as f32;
for (index, (bin, envelope)) in plan
.spectrum
.iter_mut()
.zip(plan.tail_envelope.iter_mut())
.enumerate()
{
let frequency_hz = index as f32 * bin_hz;
if frequency_hz < tuning.tail_band_low_hz || frequency_hz > tuning.tail_band_high_hz {
*envelope = *envelope * 0.70 + bin.norm() * 0.30;
continue;
}
let magnitude = bin.norm();
let expected_tail = *envelope * tuning.tail_decay;
let tail_ratio = if expected_tail > 1e-5 && magnitude < expected_tail {
((expected_tail - magnitude) / expected_tail).clamp(0.0, 1.0)
} else {
0.0
};
let reduction = (self.current_strength * tail_ratio).clamp(0.0, 0.92);
if magnitude > 1e-6 {
*bin *= 1.0 - reduction;
}
let tracking = if magnitude > *envelope { 0.32 } else { 0.12 };
*envelope = *envelope * (1.0 - tracking) + magnitude * tracking;
}
plan.inverse
.process(&mut plan.spectrum, &mut plan.output)
.map_err(|err| {
AudioError::new(format!("Tail suppression 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 tail_suppression_engages_on_midband_tail_like_signal() {
let spec = test_spec();
let mut processor = TailSuppressionProcessor::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 * 240.0 * t).sin() * 0.10;
let tail = (std::f32::consts::TAU * 2_300.0 * t).sin() * 0.045;
body + tail
})
.collect::<Vec<_>>();
let mut processed = source.clone();
for _ in 0..8 {
let mut frame = AudioFrame {
samples: source.clone(),
spec,
};
processor.process(&mut frame).unwrap();
processed = frame.samples;
}
assert!(processor.current_strength > 0.01);
assert_eq!(processed.len(), source.len());
}
}