feat: add audio visualizer implementation

src/main.rs

@@ -1,7 +1,11 @@
 use std::{
-  io::{self, stdout}, sync::{Arc, Mutex, mpsc}, thread, time::Duration,
+  io::{self, stdout},
+  sync::{Arc, Mutex, mpsc},
+  thread,
+  time::Duration,
 };
 
+use cpal::traits::{DeviceTrait, HostTrait, StreamTrait};
 use crossterm::{
   ExecutableCommand,
   event::{self, Event, KeyCode, KeyModifiers},
@@ -11,12 +15,13 @@ use crossterm::{
 use figlet_rs::Toilet;
 use notify_rust::Notification;
 use ratatui::{
-  Terminal,
+  Frame, Terminal,
   backend::CrosstermBackend,
-  layout::{Alignment, Constraint, Direction, Layout},
+  layout::{Alignment, Constraint, Direction, Layout, Rect},
   style::{Color, Style},
   widgets::{Block, Borders, Clear, Paragraph},
 };
+use realfft::{RealFftPlanner, num_complex};
 
 const SECONDS_PER_MIN: u32 = 60;
 
@@ -138,10 +143,310 @@ impl Drop for TerminalGuard {
   }
 }
 
+pub struct AudioVisualizer {
+  pub bar_data: Arc<Mutex<Vec<u64>>>,
+}
+
+// This is a implementation largely lifted from these
+// open source implementation. I just wrote it in Rust here:
+// 1. CAVA (C): https://github.com/karlstav/cava
+// 2. cli-visualizer (C++): https://github.com/dpayne/cli-visualizer
+impl AudioVisualizer {
+  pub fn new(num_bars: usize) -> Self {
+    // 
+    let bar_data = Arc::new(Mutex::new(vec![0; num_bars]));
+    let bar_data_clone = Arc::clone(&bar_data);
+
+    thread::spawn(move || {
+      if let Err(e) = Self::run_audio_loop(bar_data_clone, num_bars) {
+        eprintln!("Audio capture error: {:?}", e);
+      }
+    });
+
+    Self { bar_data }
+  }
+
+  // Doing best effort to support cross-platform functionality
+  fn get_audio_device(
+    host: &cpal::Host,
+  ) -> Result<(cpal::Device, cpal::StreamConfig), Box<dyn std::error::Error>> {
+    // Try windows WASAPI loopback on default output device
+    #[cfg(target_os = "windows")]
+    {
+      if let Some(device) = host.default_output_device()
+        && let Ok(config) = device.default_output_config()
+      {
+        return Ok((device, config.into()));
+      }
+    }
+
+    // Try Linux PipeWire/PulseAudio output monitor device
+    #[cfg(target_os = "linux")]
+    {
+      if let Ok(devices) = host.devices() {
+        for dev in devices {
+          if let Ok(name) = dev.name() {
+            // Monitor devices mirror system output under PulseAudio/PipeWire
+            if name.contains("monitor") {
+              if let Ok(config) = dev.default_input_config() {
+                return Ok((dev, config.into()));
+              }
+            }
+          }
+        }
+      }
+    }
+
+    // Fallback to default input device
+    let device = host
+      .default_input_device()
+      .or_else(|| host.default_output_device())
+      .ok_or("No audio input or output device found")?;
+
+    let config = device
+      .default_input_config()
+      .or_else(|_| device.default_output_config())?
+      .into();
+
+    Ok((device, config))
+  }
+
+  // we here in doubling octaves but fft outputs linearly spaced bins.
+  // so we bundle bins on a log scale from 20 hz to 12 kHz
+  // this allows bass, mids, and trembls to have equal visual
+  // proportions during display
+  //
+  // i.e. 
+  // bins like this [0-500Hz] [500-1k] [1k-1.5k] [1.5k-2k] [2k-2.5k] [2.5k-3k] [3k-12kHz]
+  //   vs
+  // bins like this [20-60Hz] [60-250Hz] [250-500Hz] [500-2kHz] [2k-4kHz] [4k-8kHz] [8k-12kHz]
+  fn build_log_bins(num_bars: usize, sample_rate: f32, chunk_size: usize) -> Vec<(usize, usize)> {
+    let nyquist = sample_rate / 2.0;
+    let max_hz = 12000.0f32;
+
+    (0..num_bars)
+      .map(|i| {
+        let low_hz = 20.0 * (max_hz / 20.0).powf(i as f32 / num_bars as f32);
+        let high_hz = 20.0 * (max_hz / 20.0).powf((i + 1) as f32 / num_bars as f32);
+
+        let low = ((low_hz / nyquist) * (chunk_size as f32 / 2.0)) as usize;
+        let high = ((high_hz / nyquist) * (chunk_size as f32 / 2.0)) as usize;
+
+        (low.max(1), high.max(low + 1))
+      })
+      .collect()
+  }
+
+  // without we get sharp edges which gives noise in FFT processing
+  fn apply_hann_window(samples: &[f32], input_buffer: &mut [f32]) {
+    let chunk_size = samples.len();
+
+    for (i, sample) in samples.iter().enumerate() {
+      let window = 0.5 * (1.0 - (2.0 * std::f32::consts::PI * i as f32 / chunk_size as f32).cos());
+      input_buffer[i] = sample * window;
+    }
+  }
+
+  fn process_fft_magnitudes(
+    spectrum: &[num_complex::Complex32],
+    log_bins: &[(usize, usize)],
+    freq_boost: &[f32],
+    prev_heights: &[f32],
+    autosens: f32,
+    smoothing: f32,
+    falloff: f32,
+  ) -> Vec<f32> {
+    let num_bars = log_bins.len();
+    let mut current_bars = vec![0.0f32; num_bars];
+
+    // determine magnitude
+    for i in 0..num_bars {
+      let (start, stop) = log_bins[i];
+      let bin_slice = &spectrum[start..stop.min(spectrum.len())];
+      let magnitude_sum: f32 = bin_slice.iter().map(|c| c.norm()).sum();
+      let avg_mag = if !bin_slice.is_empty() {
+        magnitude_sum / bin_slice.len() as f32
+      } else {
+        0.0
+      };
+
+      // ignore quiet static noise below specific amplitude
+      let raw_val = if avg_mag < 0.02 {
+        0.0
+      } else {
+        (avg_mag * freq_boost[i] * autosens + 1.0).log10() * 3.5
+      };
+
+      // rise smoothly from previous height
+      let target = (raw_val * (1.0 - smoothing)) + (prev_heights[i] * smoothing);
+
+      // prevent sudden drops
+      if target < prev_heights[i] {
+        current_bars[i] = (prev_heights[i] - falloff).max(0.0);
+      } else {
+        current_bars[i] = target;
+      }
+    }
+
+    current_bars
+  }
+
+  // blend the heights between neighboring freq bins
+  // so instead of sharp spikes we get more of waves
+  fn apply_monstercat_smoothing(bars: &[f32]) -> Vec<f32> {
+    let num_bars = bars.len();
+    let mut smoothed = bars.to_vec();
+
+    for i in 1..(num_bars - 1) {
+      smoothed[i] = (bars[i - 1] * 0.25) + (bars[i] * 0.50) + (bars[i + 1] * 0.25);
+    }
+
+    smoothed
+  }
+
+  // songs can be quiet and load so we try to adjust
+  // sensitivity to avoid bars becoming flattened or
+  // clipped
+  fn adjust_autosens(autosens: &mut f32, bars: &[f32]) {
+    let max_val = bars.iter().cloned().fold(0.0f32, f32::max);
+    if max_val > 8.0 {
+      *autosens *= 0.98;
+    } else if max_val < 3.0 && *autosens < 3.0 {
+      *autosens *= 1.01;
+    }
+  }
+
+  // captures output and runs through pipeline
+  // system audio -> audio buffer -> windowing -> fft process -> binning -> floor/autosens -> smoothing
+  fn run_audio_loop(
+    bar_data: Arc<Mutex<Vec<u64>>>,
+    num_bars: usize,
+  ) -> Result<(), Box<dyn std::error::Error>> {
+    let host = cpal::default_host();
+    let (device, config) = Self::get_audio_device(&host)?;
+
+    let sample_rate = config.sample_rate as f32;
+    let channels = config.channels as usize;
+
+    let chunk_size = 2048;
+    let mut planner = RealFftPlanner::<f32>::new();
+    let fft = planner.plan_fft_forward(chunk_size);
+    let mut windowed_buffer = fft.make_input_vec();
+    let mut spectrum = fft.make_output_vec();
+
+    let mut prev_heights = vec![0.0f32; num_bars];
+    let smoothing = 0.70f32;
+    let falloff = 0.08f32;
+    let mut autosens = 1.0f32;
+
+    let log_bins = Self::build_log_bins(num_bars, sample_rate, chunk_size);
+
+    // high freq have less amplitude so we boost
+    // more and more as we go right
+    let freq_boost: Vec<f32> = (0..num_bars)
+      .map(|i| 1.0 + (3.5 * (i as f32 / num_bars as f32).powf(1.2)))
+      .collect();
+
+    let audio_buffer = Arc::new(Mutex::new(Vec::<f32>::with_capacity(chunk_size * 2)));
+    let buffer_clone = Arc::clone(&audio_buffer);
+
+    let stream = device.build_input_stream(
+      config,
+      move |data: &[f32], _| {
+        let mut buf = buffer_clone.lock().unwrap();
+
+        for chunk in data.chunks(channels) {
+          let mono: f32 = chunk.iter().sum::<f32>() / channels as f32;
+          buf.push(mono);
+        }
+
+        if buf.len() > chunk_size * 2 {
+          let drain_amt = buf.len() - chunk_size;
+          buf.drain(0..drain_amt);
+        }
+      },
+      |err| eprintln!("Stream error: {}", err),
+      None,
+    )?;
+
+    stream.play()?;
+
+    loop {
+      thread::sleep(std::time::Duration::from_millis(16));
+
+      let samples = {
+        let buf = audio_buffer.lock().unwrap();
+        if buf.len() < chunk_size {
+          continue;
+        }
+        buf[buf.len() - chunk_size..].to_vec()
+      };
+
+      Self::apply_hann_window(&samples, &mut windowed_buffer);
+
+      fft.process(&mut windowed_buffer, &mut spectrum).unwrap();
+
+      let raw_bars = Self::process_fft_magnitudes(
+        &spectrum,
+        &log_bins,
+        &freq_boost,
+        &prev_heights,
+        autosens,
+        smoothing,
+        falloff,
+      );
+
+      let smoothed_bars = Self::apply_monstercat_smoothing(&raw_bars);
+      prev_heights = smoothed_bars.clone();
+
+      Self::adjust_autosens(&mut autosens, &smoothed_bars);
+
+      if let Ok(mut bars) = bar_data.lock() {
+        for (i, val) in smoothed_bars.iter().enumerate() {
+          bars[i] = ((*val * 10.0).clamp(0.0, 100.0)) as u64;
+        }
+      }
+    }
+  }
+}
+
+// use unicode blocks for rendering
+// Empty        -> " "
+// 1/8th height -> "▂"
+// 2/8th height -> "▃"
+// 3/8th height -> "▄"
+// 4/8th height -> "▅"
+// 5/8th height -> "▆"
+// 6/8th height -> "▇"
+// Full Height  -> "█"
+fn render_visualizer(f: &mut Frame, area: Rect, bar_values: &[u64]) {
+  const BLOCKS: [&str; 8] = ["▂", "▃", "▄", "▅", "▆", "▇", "█", "█"];
+
+  let line: String = bar_values
+    .iter()
+    .map(|&val| {
+      if val == 0 {
+        " "
+      } else {
+        let idx = ((val as f32 / 100.0) * (BLOCKS.len() - 1) as f32)
+          .clamp(0.0, (BLOCKS.len() - 1) as f32) as usize;
+
+        BLOCKS[idx]
+      }
+    })
+    .collect();
+
+  let viz_paragraph = Paragraph::new(line)
+    .style(Style::default().fg(Color::Cyan))
+    .alignment(Alignment::Center);
+
+  f.render_widget(viz_paragraph, area);
+}
+
 // TODO: Refactor to an Elm like architecture
 fn main() -> Result<(), io::Error> {
   let args: Vec<String> = std::env::args().collect();
-  
+
   if args.iter().any(|arg| arg == "-h" || arg == "--help") {
     println!("Usage: termato [options] [work_minutes] [break_minutes]");
     println!();
@@ -169,18 +474,31 @@ fn main() -> Result<(), io::Error> {
   let enable_visualizer = args.iter().any(|arg| arg == "-z" || arg == "--visualizer");
   let enable_notifications = args.iter().any(|arg| arg == "-n" || arg == "--notify");
 
-  let positional_args: Vec<&String> = args.iter()
+  let visualizer = if enable_visualizer {
+    Some(AudioVisualizer::new(64))
+  } else {
+    None
+  };
+
+  let positional_args: Vec<&String> = args
+    .iter()
     .skip(1)
-    .filter(|arg| { 
-        *arg != "-n" && *arg != "--notify" && *arg != "-vx" && *arg != "--visualizer"
-    })
+    .filter(|arg| *arg != "-n" && *arg != "--notify" && *arg != "-z" && *arg != "--visualizer")
     .collect();
 
-  let work_mins = positional_args.first().and_then(|s| s.parse().ok()).unwrap_or(25);
-  let break_mins = positional_args.get(1).and_then(|s| s.parse().ok()).unwrap_or(5);
+  let work_mins = positional_args
+    .first()
+    .and_then(|s| s.parse().ok())
+    .unwrap_or(25);
+
+  let break_mins = positional_args
+    .get(1)
+    .and_then(|s| s.parse().ok())
+    .unwrap_or(5);
+
   let mut termato = Termato::new(work_mins, break_mins)
-      .with_notifications(enable_notifications)
-      .with_visualizer(enable_visualizer);
+    .with_notifications(enable_notifications)
+    .with_visualizer(enable_visualizer);
 
   let (tx, rx) = mpsc::channel();
 
@@ -210,28 +528,58 @@ fn main() -> Result<(), io::Error> {
         }
       };
 
-      let big_time_text = if let Some(fig) = font.convert(&time_str) {
+      let raw_time_text = if let Some(fig) = font.convert(&time_str) {
         fig.to_string()
       } else {
         time_str
       };
 
-      let text_height = big_time_text.lines().count() as u16;
+      let trimmed_time_text = raw_time_text.trim_matches('\n');
+      let text_lines: Vec<&str> = trimmed_time_text.lines().collect();
+      let text_height = text_lines.len() as u16;
 
-      let vertical_layout = Layout::default()
+      let viz_height = if visualizer.is_some() { 1 } else { 0 };
+      let viz_width = 64;
+
+      let content_height = text_height + viz_height;
+
+      let outer_vertical = Layout::default()
         .direction(Direction::Vertical)
         .constraints([
           Constraint::Min(0),
-          Constraint::Length(text_height),
+          Constraint::Length(content_height),
           Constraint::Min(0),
         ])
         .split(size);
 
-      let timer_paragraph = Paragraph::new(big_time_text)
+      let inner_vertical = Layout::default()
+        .direction(Direction::Vertical)
+        .constraints([
+          Constraint::Length(text_height),
+          Constraint::Length(viz_height),
+        ])
+        .split(outer_vertical[1]);
+
+      let viz_horizontal = Layout::default()
+        .direction(Direction::Horizontal)
+        .constraints([
+          Constraint::Min(0),
+          Constraint::Length(viz_width.min(size.width)),
+          Constraint::Min(0),
+        ])
+        .split(inner_vertical[1]);
+
+      let timer_paragraph = Paragraph::new(trimmed_time_text)
         .style(Style::default().fg(time_color))
         .alignment(Alignment::Center);
 
-      f.render_widget(timer_paragraph, vertical_layout[1]);
+      f.render_widget(timer_paragraph, inner_vertical[0]);
+
+      if let Some(ref viz) = visualizer
+        && let Ok(bar_data) = viz.bar_data.lock()
+      {
+        render_visualizer(f, viz_horizontal[1], &bar_data);
+      }
 
       if termato.show_help {
         let popup_width = 44;
@@ -280,7 +628,7 @@ fn main() -> Result<(), io::Error> {
       }
     })?;
 
-    if event::poll(Duration::from_millis(100))?
+    if event::poll(Duration::from_millis(16))?
       && let Event::Key(key) = event::read()?
       && key.kind == event::KeyEventKind::Press
     {