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 beat_avg: f32,
252 beat_energy: f32,
254}
255
256impl AnalysisState {
257 fn new(
258 scale: FrequencyScale,
259 bar_half_life: f32,
260 peak_half_life: f32,
261 amplitude_scale: AmplitudeScale,
262 ) -> Self {
263 let mut planner = RealFftPlanner::<f32>::new();
264 let fft = planner.plan_fft_forward(FFT_SIZE);
265 let fft_input = fft.make_input_vec();
266 let fft_output = fft.make_output_vec();
267 Self {
268 window: hann_window(),
269 fft_input,
270 fft_output,
271 fft,
272 bin_to_bar: Vec::new(),
273 last_sample_rate: 0.0,
274 bar_counts: [0u32; NUM_BARS],
275 prev_spectrum: [0.0; NUM_BARS],
276 spectrum: [0.0; NUM_BARS],
277 peaks: [0.0; NUM_BARS],
278 vu_levels: [0.0; 2],
279 last_update: Instant::now(),
280 scale,
281 bar_half_life,
282 peak_half_life,
283 amplitude_scale,
284 a_weight_table: Vec::new(),
285 beat_avg: 0.0,
286 beat_energy: 0.0,
287 }
288 }
289
290 fn decay_factors(&mut self) -> (f32, f32) {
292 let now = Instant::now();
293 let dt = now.duration_since(self.last_update).as_secs_f32();
294 self.last_update = now;
295 let bar_decay = 0.5f32.powf(dt / self.bar_half_life);
296 let peak_decay = 0.5f32.powf(dt / self.peak_half_life);
297 (bar_decay, peak_decay)
298 }
299
300 fn analyze(&mut self, samples: &[f32], channels: usize, sample_rate: f32) {
304 if samples.is_empty() || sample_rate <= 0.0 || channels == 0 {
305 self.decay_silence();
306 return;
307 }
308
309 self.compute_vu(samples, channels);
311
312 let total_frames = samples.len() / channels;
314 let frames_to_use = total_frames.min(FFT_SIZE);
315 let frame_start = total_frames - frames_to_use;
316
317 for i in 0..FFT_SIZE {
318 if i < frames_to_use {
319 let frame_idx = frame_start + i;
320 let sample_start = frame_idx * channels;
321 let mut sum = 0.0f32;
322 for ch in 0..channels {
323 if sample_start + ch < samples.len() {
324 sum += samples[sample_start + ch];
325 }
326 }
327 self.fft_input[i] = (sum / channels as f32) * self.window[i];
328 } else {
329 self.fft_input[i] = 0.0;
330 }
331 }
332
333 if self
335 .fft
336 .process(&mut self.fft_input, &mut self.fft_output)
337 .is_err()
338 {
339 self.decay_silence();
340 return;
341 }
342
343 if (sample_rate - self.last_sample_rate).abs() > 0.5 {
345 self.bin_to_bar = build_bin_to_bar(sample_rate, self.scale);
346 self.a_weight_table = build_a_weight_table(sample_rate);
347 self.last_sample_rate = sample_rate;
348 }
349
350 std::mem::swap(&mut self.spectrum, &mut self.prev_spectrum);
352 for bar in self.spectrum.iter_mut() {
353 *bar = 0.0;
354 }
355 for c in self.bar_counts.iter_mut() {
356 *c = 0;
357 }
358
359 let norm = 2.0 / FFT_SIZE as f32;
360 let db_range_inv = 1.0 / (DB_CEIL - DB_FLOOR);
361 let num_bins = self.fft_output.len().min(self.bin_to_bar.len());
362
363 for bin_idx in 0..num_bins {
364 let bar_idx = match self.bin_to_bar[bin_idx] {
365 Some(b) => b,
366 None => continue,
367 };
368 let c = self.fft_output[bin_idx];
369 let magnitude = (c.re * c.re + c.im * c.im).sqrt() * norm;
370 let mut db = if magnitude > 0.0 {
371 20.0 * magnitude.log10()
372 } else {
373 DB_FLOOR
374 };
375 if matches!(
377 self.amplitude_scale,
378 AmplitudeScale::Perceptual | AmplitudeScale::AWeight
379 ) && let Some(&aw) = self.a_weight_table.get(bin_idx)
380 {
381 db += aw;
382 }
383 let level = ((db - DB_FLOOR) * db_range_inv).clamp(0.0, 1.0);
384 let level = self.amplitude_scale.apply(level);
385 if level > self.spectrum[bar_idx] {
386 self.spectrum[bar_idx] = level;
387 }
388 self.bar_counts[bar_idx] += 1;
389 }
390
391 for i in 0..NUM_BARS {
393 if self.bar_counts[i] == 0 {
394 let left = if i > 0 { self.spectrum[i - 1] } else { 0.0 };
395 let right = if i + 1 < NUM_BARS {
396 self.spectrum[i + 1]
397 } else {
398 0.0
399 };
400 self.spectrum[i] = (left + right) * 0.5;
401 }
402 }
403
404 let (bar_decay, peak_decay) = self.decay_factors();
406 for i in 0..NUM_BARS {
407 let decayed = self.prev_spectrum[i] * bar_decay;
408 self.spectrum[i] = self.spectrum[i].max(decayed);
409
410 if self.spectrum[i] > self.peaks[i] {
411 self.peaks[i] = self.spectrum[i];
412 } else {
413 self.peaks[i] *= peak_decay;
414 }
415 }
416
417 let beat_bands = NUM_BARS.min(6);
420 let low_energy: f32 = self.spectrum[..beat_bands].iter().sum::<f32>() / beat_bands as f32;
421
422 const BEAT_AVG_ALPHA: f32 = 0.02;
425 self.beat_avg = self.beat_avg * (1.0 - BEAT_AVG_ALPHA) + low_energy * BEAT_AVG_ALPHA;
426
427 let beat_spike = if self.beat_avg > 0.005 {
430 let excess = (low_energy - self.beat_avg).max(0.0);
431 (excess / self.beat_avg.max(0.05)).clamp(0.0, 1.0)
432 } else {
433 (low_energy * 3.0).clamp(0.0, 1.0)
435 };
436
437 self.beat_energy = beat_spike.max(self.beat_energy * bar_decay.sqrt());
440 }
441
442 fn decay_silence(&mut self) {
444 let (bar_decay, peak_decay) = self.decay_factors();
445 for i in 0..NUM_BARS {
446 self.spectrum[i] *= bar_decay;
447 self.peaks[i] *= peak_decay;
448 }
449 for v in self.vu_levels.iter_mut() {
450 *v *= bar_decay;
451 }
452 self.beat_energy *= bar_decay;
453 }
454
455 fn compute_vu(&mut self, samples: &[f32], channels: usize) {
457 let total_frames = samples.len() / channels;
458 let frames_to_use = total_frames.min(2048);
459 let frame_start = total_frames - frames_to_use;
460 let vu_channels = channels.min(2);
461 let mut sum_sq = [0.0f64; 2];
462
463 for frame in 0..frames_to_use {
464 let idx = (frame_start + frame) * channels;
465 for ch in 0..vu_channels {
466 if idx + ch < samples.len() {
467 let s = samples[idx + ch] as f64;
468 sum_sq[ch] += s * s;
469 }
470 }
471 }
472
473 let db_range = DB_CEIL - DB_FLOOR;
474 for (ch, &sq) in sum_sq.iter().enumerate().take(vu_channels) {
475 let rms = (sq / frames_to_use as f64).sqrt() as f32;
476 let db = if rms > 0.0 {
477 20.0 * rms.log10()
478 } else {
479 DB_FLOOR
480 };
481 self.vu_levels[ch] = ((db - DB_FLOOR) / db_range).clamp(0.0, 1.0);
482 }
483
484 if vu_channels == 1 {
485 self.vu_levels[1] = self.vu_levels[0];
486 }
487 }
488}
489
490pub struct VizAnalyzer {
501 output: SharedAnalysisOutput,
502 running: Arc<AtomicBool>,
503 handle: Option<thread::JoinHandle<()>>,
504}
505
506impl VizAnalyzer {
507 pub fn spawn(viz_buffer: Arc<VizBuffer>, cfg: &VisualizerConfig) -> Self {
512 Self::spawn_inner(viz_buffer, cfg, None)
513 }
514
515 pub fn spawn_with_snapshot(
520 viz_buffer: Arc<VizBuffer>,
521 cfg: &VisualizerConfig,
522 snapshot: Arc<VizSnapshot>,
523 ) -> Self {
524 Self::spawn_inner(viz_buffer, cfg, Some(snapshot))
525 }
526
527 fn spawn_inner(
528 viz_buffer: Arc<VizBuffer>,
529 cfg: &VisualizerConfig,
530 snapshot: Option<Arc<VizSnapshot>>,
531 ) -> Self {
532 let output: SharedAnalysisOutput = Arc::new(Mutex::new(AnalysisOutput::default()));
533 let running = Arc::new(AtomicBool::new(true));
534
535 let scale = FrequencyScale::parse(&cfg.scale);
536 let amplitude_scale = AmplitudeScale::parse(&cfg.amplitude_scale);
537 let bar_half_life = cfg.bar_decay_ms as f32 / 1000.0;
538 let peak_half_life = cfg.peak_decay_ms as f32 / 1000.0;
539 let interval = Duration::from_millis(1000 / cfg.fps.max(1) as u64);
540
541 let output_clone = Arc::clone(&output);
542 let running_clone = Arc::clone(&running);
543
544 let handle = thread::Builder::new()
545 .name("viz-analyzer".into())
546 .spawn(move || {
547 analysis_loop(
548 viz_buffer,
549 output_clone,
550 snapshot,
551 running_clone,
552 scale,
553 amplitude_scale,
554 bar_half_life,
555 peak_half_life,
556 interval,
557 );
558 })
559 .expect("failed to spawn viz-analyzer thread");
560
561 Self {
562 output,
563 running,
564 handle: Some(handle),
565 }
566 }
567
568 pub fn output(&self) -> AnalysisOutput {
573 self.output.lock().clone()
574 }
575
576 pub fn shared_output(&self) -> SharedAnalysisOutput {
579 Arc::clone(&self.output)
580 }
581
582 pub fn shutdown(&mut self) {
584 self.running.store(false, Ordering::Relaxed);
585 if let Some(h) = self.handle.take() {
586 let _ = h.join();
587 }
588 }
589}
590
591impl Drop for VizAnalyzer {
592 fn drop(&mut self) {
593 self.shutdown();
594 }
595}
596
597#[allow(clippy::too_many_arguments)]
600fn analysis_loop(
601 viz_buffer: Arc<VizBuffer>,
602 output: SharedAnalysisOutput,
603 snapshot: Option<Arc<VizSnapshot>>,
604 running: Arc<AtomicBool>,
605 scale: FrequencyScale,
606 amplitude_scale: AmplitudeScale,
607 bar_half_life: f32,
608 peak_half_life: f32,
609 interval: Duration,
610) {
611 let mut state = AnalysisState::new(scale, bar_half_life, peak_half_life, amplitude_scale);
612
613 while running.load(Ordering::Relaxed) {
614 let start = Instant::now();
615
616 let snap: RawVizSnapshot = viz_buffer.snapshot_with_meta();
618
619 state.analyze(
621 &snap.samples,
622 snap.channels.max(1) as usize,
623 snap.sample_rate as f32,
624 );
625
626 {
628 let mut out = output.lock();
629 out.spectrum.copy_from_slice(&state.spectrum);
630 out.peaks.copy_from_slice(&state.peaks);
631 out.vu_levels = state.vu_levels;
632 }
633
634 if let Some(ref snap_out) = snapshot {
636 use super::viz::WAVEFORM_SAMPLES;
639 let waveform = if snap.channels >= 1 && !snap.samples.is_empty() {
640 let interleaved_len = WAVEFORM_SAMPLES * snap.channels.max(1) as usize;
641 let start = snap.samples.len().saturating_sub(interleaved_len);
642 snap.samples[start..].to_vec()
643 } else {
644 Vec::new()
645 };
646 snap_out.write(VizFrame {
647 spectrum: state.spectrum,
648 peaks: state.peaks,
649 vu_levels: state.vu_levels,
650 beat_energy: state.beat_energy,
651 timestamp: Instant::now(),
652 waveform,
653 });
654 }
655
656 let elapsed = start.elapsed();
658 if elapsed < interval {
659 thread::sleep(interval - elapsed);
660 }
661 }
662}
663
664#[cfg(test)]
667mod tests {
668 use super::*;
669 use crate::audio::viz::VizBuffer;
670 use crate::config::VisualizerConfig;
671
672 fn make_cfg() -> VisualizerConfig {
673 VisualizerConfig::default()
674 }
675
676 #[test]
677 fn analyzer_spawns_and_shuts_down() {
678 let buf = VizBuffer::new();
679 let cfg = make_cfg();
680 let mut analyzer = VizAnalyzer::spawn(buf, &cfg);
681 std::thread::sleep(Duration::from_millis(100));
683 analyzer.shutdown();
684 let out = analyzer.output();
686 assert_eq!(out.spectrum.len(), NUM_BARS);
687 assert_eq!(out.peaks.len(), NUM_BARS);
688 }
689
690 #[test]
691 fn analyzer_produces_nonzero_output_for_sine() {
692 let buf = VizBuffer::new();
693 let sample_rate = 44100u32;
694 let channels = 2u16;
695 let num_frames = 4096;
696 let mut samples = Vec::with_capacity(num_frames * 2);
697 for i in 0..num_frames {
698 let t = i as f32 / sample_rate as f32;
699 let val = (2.0 * std::f32::consts::PI * 440.0 * t).sin() * 0.5;
700 samples.push(val);
701 samples.push(val);
702 }
703 buf.push_samples(&samples, channels, sample_rate);
704
705 let cfg = make_cfg();
706 let mut analyzer = VizAnalyzer::spawn(Arc::clone(&buf), &cfg);
707 std::thread::sleep(Duration::from_millis(150));
709
710 let out = analyzer.output();
711 analyzer.shutdown();
712
713 let max_bar = out.spectrum.iter().cloned().fold(0.0f32, f32::max);
714 assert!(
715 max_bar > 0.05,
716 "expected nonzero spectrum for 440 Hz sine, max = {}",
717 max_bar
718 );
719 }
720
721 #[test]
722 fn analysis_state_decays_to_zero_on_silence() {
723 let mut state =
726 AnalysisState::new(FrequencyScale::Bark, 0.08, 0.35, AmplitudeScale::Linear);
727
728 for v in state.spectrum.iter_mut() {
730 *v = 1.0;
731 }
732 for v in state.peaks.iter_mut() {
733 *v = 1.0;
734 }
735
736 let silence: Vec<f32> = vec![0.0; FFT_SIZE * 2];
741 for _ in 0..100 {
742 state.last_update = Instant::now() - Duration::from_millis(100);
743 state.analyze(&silence, 2, 44100.0);
744 }
745
746 let max_spec = state.spectrum.iter().cloned().fold(0.0f32, f32::max);
747 let max_peak = state.peaks.iter().cloned().fold(0.0f32, f32::max);
748 assert!(
749 max_spec < 0.1,
750 "spectrum should decay near zero, got {}",
751 max_spec
752 );
753 assert!(
754 max_peak < 0.1,
755 "peaks should decay near zero, got {}",
756 max_peak
757 );
758 }
759
760 #[test]
761 fn bin_to_bar_covers_audible_range() {
762 let mapping = build_bin_to_bar(44100.0, FrequencyScale::Bark);
763 let active_bins: Vec<usize> = mapping.iter().filter_map(|x| *x).collect();
764 assert!(
765 !active_bins.is_empty(),
766 "at least some bins should map to bars"
767 );
768 let max_bar = *active_bins.iter().max().unwrap();
769 assert!(max_bar < NUM_BARS, "bar index must be in range");
770 }
771
772 #[test]
773 fn frequency_scale_bark_normalize_monotonic() {
774 let scale = FrequencyScale::Bark;
775 let freqs: Vec<f32> = vec![100.0, 500.0, 1000.0, 4000.0, 10000.0];
776 let normed: Vec<f32> = freqs.iter().map(|&f| scale.normalize(f)).collect();
777 for w in normed.windows(2) {
778 assert!(w[1] > w[0], "Bark scale must be monotonically increasing");
779 }
780 }
781}