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 vu_levels: state.vu_levels,
649 beat_energy: state.beat_energy,
650 timestamp: Instant::now(),
651 waveform,
652 });
653 }
654
655 let elapsed = start.elapsed();
657 if elapsed < interval {
658 thread::sleep(interval - elapsed);
659 }
660 }
661}
662
663#[cfg(test)]
666mod tests {
667 use super::*;
668 use crate::audio::viz::VizBuffer;
669 use crate::config::VisualizerConfig;
670
671 fn make_cfg() -> VisualizerConfig {
672 VisualizerConfig::default()
673 }
674
675 #[test]
676 fn analyzer_spawns_and_shuts_down() {
677 let buf = VizBuffer::new();
678 let cfg = make_cfg();
679 let mut analyzer = VizAnalyzer::spawn(buf, &cfg);
680 std::thread::sleep(Duration::from_millis(100));
682 analyzer.shutdown();
683 let out = analyzer.output();
685 assert_eq!(out.spectrum.len(), NUM_BARS);
686 assert_eq!(out.peaks.len(), NUM_BARS);
687 }
688
689 #[test]
690 fn analyzer_produces_nonzero_output_for_sine() {
691 let buf = VizBuffer::new();
692 let sample_rate = 44100u32;
693 let channels = 2u16;
694 let num_frames = 4096;
695 let mut samples = Vec::with_capacity(num_frames * 2);
696 for i in 0..num_frames {
697 let t = i as f32 / sample_rate as f32;
698 let val = (2.0 * std::f32::consts::PI * 440.0 * t).sin() * 0.5;
699 samples.push(val);
700 samples.push(val);
701 }
702 buf.push_samples(&samples, channels, sample_rate);
703
704 let cfg = make_cfg();
705 let mut analyzer = VizAnalyzer::spawn(Arc::clone(&buf), &cfg);
706 std::thread::sleep(Duration::from_millis(150));
708
709 let out = analyzer.output();
710 analyzer.shutdown();
711
712 let max_bar = out.spectrum.iter().cloned().fold(0.0f32, f32::max);
713 assert!(
714 max_bar > 0.05,
715 "expected nonzero spectrum for 440 Hz sine, max = {}",
716 max_bar
717 );
718 }
719
720 #[test]
721 fn analysis_state_decays_to_zero_on_silence() {
722 let mut state =
725 AnalysisState::new(FrequencyScale::Bark, 0.08, 0.35, AmplitudeScale::Linear);
726
727 for v in state.spectrum.iter_mut() {
729 *v = 1.0;
730 }
731 for v in state.peaks.iter_mut() {
732 *v = 1.0;
733 }
734
735 let silence: Vec<f32> = vec![0.0; FFT_SIZE * 2];
740 for _ in 0..100 {
741 state.last_update = Instant::now() - Duration::from_millis(100);
742 state.analyze(&silence, 2, 44100.0);
743 }
744
745 let max_spec = state.spectrum.iter().cloned().fold(0.0f32, f32::max);
746 let max_peak = state.peaks.iter().cloned().fold(0.0f32, f32::max);
747 assert!(
748 max_spec < 0.1,
749 "spectrum should decay near zero, got {}",
750 max_spec
751 );
752 assert!(
753 max_peak < 0.1,
754 "peaks should decay near zero, got {}",
755 max_peak
756 );
757 }
758
759 #[test]
760 fn bin_to_bar_covers_audible_range() {
761 let mapping = build_bin_to_bar(44100.0, FrequencyScale::Bark);
762 let active_bins: Vec<usize> = mapping.iter().filter_map(|x| *x).collect();
763 assert!(
764 !active_bins.is_empty(),
765 "at least some bins should map to bars"
766 );
767 let max_bar = *active_bins.iter().max().unwrap();
768 assert!(max_bar < NUM_BARS, "bar index must be in range");
769 }
770
771 #[test]
772 fn frequency_scale_bark_normalize_monotonic() {
773 let scale = FrequencyScale::Bark;
774 let freqs: Vec<f32> = vec![100.0, 500.0, 1000.0, 4000.0, 10000.0];
775 let normed: Vec<f32> = freqs.iter().map(|&f| scale.normalize(f)).collect();
776 for w in normed.windows(2) {
777 assert!(w[1] > w[0], "Bark scale must be monotonically increasing");
778 }
779 }
780}