1use std::sync::Arc;
21use std::sync::atomic::{AtomicBool, Ordering};
22use std::thread;
23use std::time::{Duration, Instant};
24
25use parking_lot::Mutex;
26use realfft::RealFftPlanner;
27
28use super::viz::{
29 AnalysisOutput, NUM_BARS, RawVizSnapshot, SharedAnalysisOutput, VizBuffer, VizFrame,
30 VizSnapshot,
31};
32use crate::config::VisualizerConfig;
33
34const FFT_SIZE: usize = 2048;
38
39const MIN_FREQ: f32 = 20.0;
41
42const MAX_FREQ: f32 = 18_000.0;
44
45const DB_FLOOR: f32 = -80.0;
47
48const DB_CEIL: f32 = 0.0;
50
51#[derive(Debug, Clone, Copy, Default)]
55pub enum FrequencyScale {
56 #[default]
58 Bark,
59 Mel,
61 Log,
63 Linear,
65}
66
67impl FrequencyScale {
68 pub fn parse(s: &str) -> Self {
69 match s.to_lowercase().as_str() {
70 "bark" => Self::Bark,
71 "mel" => Self::Mel,
72 "log" | "logarithmic" => Self::Log,
73 "linear" => Self::Linear,
74 _ => Self::default(),
75 }
76 }
77
78 fn normalize(&self, freq: f32) -> f32 {
80 match self {
81 Self::Bark => {
82 let bark = |f: f32| 26.81 / (1.0 + 1960.0 / f) - 0.53;
83 let b = bark(freq);
84 let b_min = bark(MIN_FREQ);
85 let b_max = bark(MAX_FREQ);
86 (b - b_min) / (b_max - b_min)
87 }
88 Self::Mel => {
89 let mel = |f: f32| 2595.0 * (1.0 + f / 700.0).log10();
90 let m = mel(freq);
91 let m_min = mel(MIN_FREQ);
92 let m_max = mel(MAX_FREQ);
93 (m - m_min) / (m_max - m_min)
94 }
95 Self::Log => {
96 let log_min = MIN_FREQ.ln();
97 let log_max = MAX_FREQ.ln();
98 (freq.ln() - log_min) / (log_max - log_min)
99 }
100 Self::Linear => (freq - MIN_FREQ) / (MAX_FREQ - MIN_FREQ),
101 }
102 }
103}
104
105#[derive(Debug, Clone, Copy, Default)]
109pub enum AmplitudeScale {
110 Perceptual,
112 #[default]
114 AWeight,
115 Sqrt,
117 Linear,
119}
120
121impl AmplitudeScale {
122 pub fn parse(s: &str) -> Self {
123 match s.to_lowercase().as_str() {
124 "perceptual" => Self::Perceptual,
125 "aweight" | "a-weight" | "a_weight" => Self::AWeight,
126 "sqrt" => Self::Sqrt,
127 "linear" => Self::Linear,
128 _ => Self::default(),
129 }
130 }
131
132 fn apply(self, level: f32) -> f32 {
134 match self {
135 Self::Perceptual => level.powf(0.4),
136 Self::AWeight => level,
137 Self::Sqrt => level.sqrt(),
138 Self::Linear => level,
139 }
140 }
141}
142
143fn a_weight_db(freq: f32) -> f32 {
148 let f2 = freq * freq;
149 let f4 = f2 * f2;
150
151 let num = 12194.0_f32.powi(2) * f4;
152 let denom = (f2 + 20.6_f32.powi(2))
153 * ((f2 + 107.7_f32.powi(2)) * (f2 + 737.9_f32.powi(2))).sqrt()
154 * (f2 + 12194.0_f32.powi(2));
155
156 if denom == 0.0 {
157 return DB_FLOOR;
158 }
159
160 let ra = num / denom;
162 20.0 * ra.log10() + 2.0
164}
165
166fn build_a_weight_table(sample_rate: f32) -> Vec<f32> {
168 let bin_hz = sample_rate / FFT_SIZE as f32;
169 let num_bins = FFT_SIZE / 2 + 1;
170 (0..num_bins)
171 .map(|bin_idx| {
172 let freq = bin_idx as f32 * bin_hz;
173 if freq < 1.0 {
174 DB_FLOOR } else {
176 a_weight_db(freq)
177 }
178 })
179 .collect()
180}
181
182fn hann_window() -> Vec<f32> {
186 (0..FFT_SIZE)
187 .map(|i| {
188 let t = std::f32::consts::PI * 2.0 * i as f32 / FFT_SIZE as f32;
189 0.5 * (1.0 - t.cos())
190 })
191 .collect()
192}
193
194fn build_bin_to_bar(sample_rate: f32, scale: FrequencyScale) -> Vec<Option<usize>> {
197 let bin_hz = sample_rate / FFT_SIZE as f32;
198 let num_bins = FFT_SIZE / 2 + 1;
199 (0..num_bins)
200 .map(|bin_idx| {
201 let freq = bin_idx as f32 * bin_hz;
202 if !(MIN_FREQ..=MAX_FREQ).contains(&freq) {
203 return None;
204 }
205 let normalized = scale.normalize(freq);
206 Some(((normalized * NUM_BARS as f32) as usize).min(NUM_BARS - 1))
207 })
208 .collect()
209}
210
211struct AnalysisState {
215 window: Vec<f32>,
217 fft_input: Vec<f32>,
219 fft_output: Vec<realfft::num_complex::Complex<f32>>,
221 fft: Arc<dyn realfft::RealToComplex<f32>>,
223 bin_to_bar: Vec<Option<usize>>,
225 last_sample_rate: f32,
227 bar_counts: [u32; NUM_BARS],
229 prev_spectrum: [f32; NUM_BARS],
231 spectrum: [f32; NUM_BARS],
233 peaks: [f32; NUM_BARS],
235 vu_levels: [f32; 2],
237 last_update: Instant,
239 scale: FrequencyScale,
241 bar_half_life: f32,
243 peak_half_life: f32,
245 amplitude_scale: AmplitudeScale,
247 a_weight_table: Vec<f32>,
249}
250
251impl AnalysisState {
252 fn new(
253 scale: FrequencyScale,
254 bar_half_life: f32,
255 peak_half_life: f32,
256 amplitude_scale: AmplitudeScale,
257 ) -> Self {
258 let mut planner = RealFftPlanner::<f32>::new();
259 let fft = planner.plan_fft_forward(FFT_SIZE);
260 let fft_input = fft.make_input_vec();
261 let fft_output = fft.make_output_vec();
262 Self {
263 window: hann_window(),
264 fft_input,
265 fft_output,
266 fft,
267 bin_to_bar: Vec::new(),
268 last_sample_rate: 0.0,
269 bar_counts: [0u32; NUM_BARS],
270 prev_spectrum: [0.0; NUM_BARS],
271 spectrum: [0.0; NUM_BARS],
272 peaks: [0.0; NUM_BARS],
273 vu_levels: [0.0; 2],
274 last_update: Instant::now(),
275 scale,
276 bar_half_life,
277 peak_half_life,
278 amplitude_scale,
279 a_weight_table: Vec::new(),
280 }
281 }
282
283 fn decay_factors(&mut self) -> (f32, f32) {
285 let now = Instant::now();
286 let dt = now.duration_since(self.last_update).as_secs_f32();
287 self.last_update = now;
288 let bar_decay = 0.5f32.powf(dt / self.bar_half_life);
289 let peak_decay = 0.5f32.powf(dt / self.peak_half_life);
290 (bar_decay, peak_decay)
291 }
292
293 fn analyze(&mut self, samples: &[f32], channels: usize, sample_rate: f32) {
297 if samples.is_empty() || sample_rate <= 0.0 || channels == 0 {
298 self.decay_silence();
299 return;
300 }
301
302 self.compute_vu(samples, channels);
304
305 let total_frames = samples.len() / channels;
307 let frames_to_use = total_frames.min(FFT_SIZE);
308 let frame_start = total_frames - frames_to_use;
309
310 for i in 0..FFT_SIZE {
311 if i < frames_to_use {
312 let frame_idx = frame_start + i;
313 let sample_start = frame_idx * channels;
314 let mut sum = 0.0f32;
315 for ch in 0..channels {
316 if sample_start + ch < samples.len() {
317 sum += samples[sample_start + ch];
318 }
319 }
320 self.fft_input[i] = (sum / channels as f32) * self.window[i];
321 } else {
322 self.fft_input[i] = 0.0;
323 }
324 }
325
326 if self
328 .fft
329 .process(&mut self.fft_input, &mut self.fft_output)
330 .is_err()
331 {
332 self.decay_silence();
333 return;
334 }
335
336 if (sample_rate - self.last_sample_rate).abs() > 0.5 {
338 self.bin_to_bar = build_bin_to_bar(sample_rate, self.scale);
339 self.a_weight_table = build_a_weight_table(sample_rate);
340 self.last_sample_rate = sample_rate;
341 }
342
343 std::mem::swap(&mut self.spectrum, &mut self.prev_spectrum);
345 for bar in self.spectrum.iter_mut() {
346 *bar = 0.0;
347 }
348 for c in self.bar_counts.iter_mut() {
349 *c = 0;
350 }
351
352 let norm = 2.0 / FFT_SIZE as f32;
353 let db_range_inv = 1.0 / (DB_CEIL - DB_FLOOR);
354 let num_bins = self.fft_output.len().min(self.bin_to_bar.len());
355
356 for bin_idx in 0..num_bins {
357 let bar_idx = match self.bin_to_bar[bin_idx] {
358 Some(b) => b,
359 None => continue,
360 };
361 let c = self.fft_output[bin_idx];
362 let magnitude = (c.re * c.re + c.im * c.im).sqrt() * norm;
363 let mut db = if magnitude > 0.0 {
364 20.0 * magnitude.log10()
365 } else {
366 DB_FLOOR
367 };
368 if matches!(
370 self.amplitude_scale,
371 AmplitudeScale::Perceptual | AmplitudeScale::AWeight
372 ) && let Some(&aw) = self.a_weight_table.get(bin_idx)
373 {
374 db += aw;
375 }
376 let level = ((db - DB_FLOOR) * db_range_inv).clamp(0.0, 1.0);
377 let level = self.amplitude_scale.apply(level);
378 if level > self.spectrum[bar_idx] {
379 self.spectrum[bar_idx] = level;
380 }
381 self.bar_counts[bar_idx] += 1;
382 }
383
384 for i in 0..NUM_BARS {
386 if self.bar_counts[i] == 0 {
387 let left = if i > 0 { self.spectrum[i - 1] } else { 0.0 };
388 let right = if i + 1 < NUM_BARS {
389 self.spectrum[i + 1]
390 } else {
391 0.0
392 };
393 self.spectrum[i] = (left + right) * 0.5;
394 }
395 }
396
397 let (bar_decay, peak_decay) = self.decay_factors();
399 for i in 0..NUM_BARS {
400 let decayed = self.prev_spectrum[i] * bar_decay;
401 self.spectrum[i] = self.spectrum[i].max(decayed);
402
403 if self.spectrum[i] > self.peaks[i] {
404 self.peaks[i] = self.spectrum[i];
405 } else {
406 self.peaks[i] *= peak_decay;
407 }
408 }
409 }
410
411 fn decay_silence(&mut self) {
413 let (bar_decay, peak_decay) = self.decay_factors();
414 for i in 0..NUM_BARS {
415 self.spectrum[i] *= bar_decay;
416 self.peaks[i] *= peak_decay;
417 }
418 for v in self.vu_levels.iter_mut() {
419 *v *= bar_decay;
420 }
421 }
422
423 fn compute_vu(&mut self, samples: &[f32], channels: usize) {
425 let total_frames = samples.len() / channels;
426 let frames_to_use = total_frames.min(2048);
427 let frame_start = total_frames - frames_to_use;
428 let vu_channels = channels.min(2);
429 let mut sum_sq = [0.0f64; 2];
430
431 for frame in 0..frames_to_use {
432 let idx = (frame_start + frame) * channels;
433 for ch in 0..vu_channels {
434 if idx + ch < samples.len() {
435 let s = samples[idx + ch] as f64;
436 sum_sq[ch] += s * s;
437 }
438 }
439 }
440
441 let db_range = DB_CEIL - DB_FLOOR;
442 for (ch, &sq) in sum_sq.iter().enumerate().take(vu_channels) {
443 let rms = (sq / frames_to_use as f64).sqrt() as f32;
444 let db = if rms > 0.0 {
445 20.0 * rms.log10()
446 } else {
447 DB_FLOOR
448 };
449 self.vu_levels[ch] = ((db - DB_FLOOR) / db_range).clamp(0.0, 1.0);
450 }
451
452 if vu_channels == 1 {
453 self.vu_levels[1] = self.vu_levels[0];
454 }
455 }
456}
457
458pub struct VizAnalyzer {
469 output: SharedAnalysisOutput,
470 running: Arc<AtomicBool>,
471 handle: Option<thread::JoinHandle<()>>,
472}
473
474impl VizAnalyzer {
475 pub fn spawn(viz_buffer: Arc<VizBuffer>, cfg: &VisualizerConfig) -> Self {
480 Self::spawn_inner(viz_buffer, cfg, None)
481 }
482
483 pub fn spawn_with_snapshot(
488 viz_buffer: Arc<VizBuffer>,
489 cfg: &VisualizerConfig,
490 snapshot: Arc<VizSnapshot>,
491 ) -> Self {
492 Self::spawn_inner(viz_buffer, cfg, Some(snapshot))
493 }
494
495 fn spawn_inner(
496 viz_buffer: Arc<VizBuffer>,
497 cfg: &VisualizerConfig,
498 snapshot: Option<Arc<VizSnapshot>>,
499 ) -> Self {
500 let output: SharedAnalysisOutput = Arc::new(Mutex::new(AnalysisOutput::default()));
501 let running = Arc::new(AtomicBool::new(true));
502
503 let scale = FrequencyScale::parse(&cfg.scale);
504 let amplitude_scale = AmplitudeScale::parse(&cfg.amplitude_scale);
505 let bar_half_life = cfg.bar_decay_ms as f32 / 1000.0;
506 let peak_half_life = cfg.peak_decay_ms as f32 / 1000.0;
507 let interval = Duration::from_millis(1000 / cfg.fps.max(1) as u64);
508
509 let output_clone = Arc::clone(&output);
510 let running_clone = Arc::clone(&running);
511
512 let handle = thread::Builder::new()
513 .name("viz-analyzer".into())
514 .spawn(move || {
515 analysis_loop(
516 viz_buffer,
517 output_clone,
518 snapshot,
519 running_clone,
520 scale,
521 amplitude_scale,
522 bar_half_life,
523 peak_half_life,
524 interval,
525 );
526 })
527 .expect("failed to spawn viz-analyzer thread");
528
529 Self {
530 output,
531 running,
532 handle: Some(handle),
533 }
534 }
535
536 pub fn output(&self) -> AnalysisOutput {
541 self.output.lock().clone()
542 }
543
544 pub fn shared_output(&self) -> SharedAnalysisOutput {
547 Arc::clone(&self.output)
548 }
549
550 pub fn shutdown(&mut self) {
552 self.running.store(false, Ordering::Relaxed);
553 if let Some(h) = self.handle.take() {
554 let _ = h.join();
555 }
556 }
557}
558
559impl Drop for VizAnalyzer {
560 fn drop(&mut self) {
561 self.shutdown();
562 }
563}
564
565#[allow(clippy::too_many_arguments)]
568fn analysis_loop(
569 viz_buffer: Arc<VizBuffer>,
570 output: SharedAnalysisOutput,
571 snapshot: Option<Arc<VizSnapshot>>,
572 running: Arc<AtomicBool>,
573 scale: FrequencyScale,
574 amplitude_scale: AmplitudeScale,
575 bar_half_life: f32,
576 peak_half_life: f32,
577 interval: Duration,
578) {
579 let mut state = AnalysisState::new(scale, bar_half_life, peak_half_life, amplitude_scale);
580
581 while running.load(Ordering::Relaxed) {
582 let start = Instant::now();
583
584 let snap: RawVizSnapshot = viz_buffer.snapshot_with_meta();
586
587 state.analyze(
589 &snap.samples,
590 snap.channels.max(1) as usize,
591 snap.sample_rate as f32,
592 );
593
594 {
596 let mut out = output.lock();
597 out.spectrum.copy_from_slice(&state.spectrum);
598 out.peaks.copy_from_slice(&state.peaks);
599 out.vu_levels = state.vu_levels;
600 }
601
602 if let Some(ref snap_out) = snapshot {
604 snap_out.write(VizFrame {
605 spectrum: state.spectrum,
606 vu_levels: state.vu_levels,
607 timestamp: Instant::now(),
608 });
609 }
610
611 let elapsed = start.elapsed();
613 if elapsed < interval {
614 thread::sleep(interval - elapsed);
615 }
616 }
617}
618
619#[cfg(test)]
622mod tests {
623 use super::*;
624 use crate::audio::viz::VizBuffer;
625 use crate::config::VisualizerConfig;
626
627 fn make_cfg() -> VisualizerConfig {
628 VisualizerConfig::default()
629 }
630
631 #[test]
632 fn analyzer_spawns_and_shuts_down() {
633 let buf = VizBuffer::new();
634 let cfg = make_cfg();
635 let mut analyzer = VizAnalyzer::spawn(buf, &cfg);
636 std::thread::sleep(Duration::from_millis(100));
638 analyzer.shutdown();
639 let out = analyzer.output();
641 assert_eq!(out.spectrum.len(), NUM_BARS);
642 assert_eq!(out.peaks.len(), NUM_BARS);
643 }
644
645 #[test]
646 fn analyzer_produces_nonzero_output_for_sine() {
647 let buf = VizBuffer::new();
648 let sample_rate = 44100u32;
649 let channels = 2u16;
650 let num_frames = 4096;
651 let mut samples = Vec::with_capacity(num_frames * 2);
652 for i in 0..num_frames {
653 let t = i as f32 / sample_rate as f32;
654 let val = (2.0 * std::f32::consts::PI * 440.0 * t).sin() * 0.5;
655 samples.push(val);
656 samples.push(val);
657 }
658 buf.push_samples(&samples, channels, sample_rate);
659
660 let cfg = make_cfg();
661 let mut analyzer = VizAnalyzer::spawn(Arc::clone(&buf), &cfg);
662 std::thread::sleep(Duration::from_millis(150));
664
665 let out = analyzer.output();
666 analyzer.shutdown();
667
668 let max_bar = out.spectrum.iter().cloned().fold(0.0f32, f32::max);
669 assert!(
670 max_bar > 0.05,
671 "expected nonzero spectrum for 440 Hz sine, max = {}",
672 max_bar
673 );
674 }
675
676 #[test]
677 fn analysis_state_decays_to_zero_on_silence() {
678 let mut state =
681 AnalysisState::new(FrequencyScale::Bark, 0.08, 0.35, AmplitudeScale::Linear);
682
683 for v in state.spectrum.iter_mut() {
685 *v = 1.0;
686 }
687 for v in state.peaks.iter_mut() {
688 *v = 1.0;
689 }
690
691 let silence: Vec<f32> = vec![0.0; FFT_SIZE * 2];
696 for _ in 0..100 {
697 state.last_update = Instant::now() - Duration::from_millis(100);
698 state.analyze(&silence, 2, 44100.0);
699 }
700
701 let max_spec = state.spectrum.iter().cloned().fold(0.0f32, f32::max);
702 let max_peak = state.peaks.iter().cloned().fold(0.0f32, f32::max);
703 assert!(
704 max_spec < 0.1,
705 "spectrum should decay near zero, got {}",
706 max_spec
707 );
708 assert!(
709 max_peak < 0.1,
710 "peaks should decay near zero, got {}",
711 max_peak
712 );
713 }
714
715 #[test]
716 fn bin_to_bar_covers_audible_range() {
717 let mapping = build_bin_to_bar(44100.0, FrequencyScale::Bark);
718 let active_bins: Vec<usize> = mapping.iter().filter_map(|x| *x).collect();
719 assert!(
720 !active_bins.is_empty(),
721 "at least some bins should map to bars"
722 );
723 let max_bar = *active_bins.iter().max().unwrap();
724 assert!(max_bar < NUM_BARS, "bar index must be in range");
725 }
726
727 #[test]
728 fn frequency_scale_bark_normalize_monotonic() {
729 let scale = FrequencyScale::Bark;
730 let freqs: Vec<f32> = vec![100.0, 500.0, 1000.0, 4000.0, 10000.0];
731 let normed: Vec<f32> = freqs.iter().map(|&f| scale.normalize(f)).collect();
732 for w in normed.windows(2) {
733 assert!(w[1] > w[0], "Bark scale must be monotonically increasing");
734 }
735 }
736}