SoulStream/crates/player/src/waveform.rs

//! Waveform extraction: decode an audio file and downsample it to a small
//! array of normalized 0..1 loudness values, for a waveform scrubber.
//!
//! Uses RMS energy per window, not peak amplitude: modern masters are
//! limited so hard that nearly every window's *peak* touches the digital
//! ceiling, which renders as a clipped brick. RMS tracks perceived
//! loudness and keeps the verse/chorus dynamics visible.

use std::path::Path;

use symphonia::core::audio::SampleBuffer;
use symphonia::core::codecs::DecoderOptions;
use symphonia::core::formats::{FormatOptions, SeekMode, SeekTo};
use symphonia::core::io::MediaSourceStream;
use symphonia::core::meta::MetadataOptions;
use symphonia::core::probe::Hint;
use symphonia::core::units::Time;

/// Peaks folded over this many frames before resampling to the final bucket
/// count. Keeps memory bounded regardless of track length.
const WINDOW: usize = 1024;

/// Decode `path` and return `buckets` normalized peak values (0.0..=1.0).
/// Returns an empty vec on any decode failure.
pub fn compute_waveform(path: &Path, buckets: usize) -> Vec<f32> {
    match try_compute(path, buckets) {
        Ok(v) => v,
        Err(e) => {
            log::debug!("waveform for {} failed: {e}", path.display());
            Vec::new()
        }
    }
}

/// Like [`compute_waveform`], but for one SLICE of the file — the setlist
/// segment rows' lane. `end_secs: None` runs to the file's end. The decode
/// seeks coarsely and then trims by packet timestamps, so a 4-minute slice
/// of a 90-minute set costs a 4-minute decode, not a 90-minute one.
pub fn compute_waveform_range(
    path: &Path,
    buckets: usize,
    start_secs: f64,
    end_secs: Option<f64>,
) -> Vec<f32> {
    match coarse_rms_range(path, start_secs, end_secs).map(|c| resample(&c, buckets)) {
        Ok(v) => v,
        Err(e) => {
            log::debug!("range waveform for {} failed: {e}", path.display());
            Vec::new()
        }
    }
}

fn try_compute(path: &Path, buckets: usize) -> Result<Vec<f32>, String> {
    Ok(resample(&coarse_rms(path)?, buckets))
}

/// Decode `path` into one RMS value per [`WINDOW`] frames (~23ms at 44.1kHz).
///
/// Split out from [`try_compute`] so the bucketing strategy can be measured
/// against a real decode without re-reading the file per variant — see
/// `bucketing_strategies_on_a_real_track`.
fn coarse_rms(path: &Path) -> Result<Vec<f32>, String> {
    coarse_rms_range(path, 0.0, None)
}

fn coarse_rms_range(
    path: &Path,
    start_secs: f64,
    end_secs: Option<f64>,
) -> Result<Vec<f32>, String> {
    let file = std::fs::File::open(path).map_err(|e| e.to_string())?;
    let mss = MediaSourceStream::new(Box::new(file), Default::default());

    let mut hint = Hint::new();
    if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
        // Stream partials are "<name>.<ext>.partial" — hint the codec ext.
        let ext = if ext.eq_ignore_ascii_case("partial") {
            path.file_stem()
                .map(Path::new)
                .and_then(|p| p.extension())
                .and_then(|e| e.to_str())
                .unwrap_or(ext)
        } else {
            ext
        };
        hint.with_extension(ext);
    }

    let probed = symphonia::default::get_probe()
        .format(
            &hint,
            mss,
            &FormatOptions::default(),
            &MetadataOptions::default(),
        )
        .map_err(|e| e.to_string())?;
    let mut format = probed.format;

    let track = format
        .default_track()
        .ok_or_else(|| "no default track".to_string())?;
    let track_id = track.id;
    let time_base = track.codec_params.time_base;
    let sample_rate = track.codec_params.sample_rate;
    let mut decoder = symphonia::default::get_codecs()
        .make(&track.codec_params, &DecoderOptions::default())
        .map_err(|e| e.to_string())?;

    if start_secs > 0.0 {
        // A coarse seek lands at a sync point near (usually before) the
        // target; the timestamp gate in the loop trims the difference. A
        // format that cannot seek just decodes from the top — slower, same
        // result.
        let _ = format.seek(
            SeekMode::Coarse,
            SeekTo::Time {
                time: Time::from(start_secs),
                track_id: Some(track_id),
            },
        );
        decoder.reset();
    }

    let mut sample_buf: Option<SampleBuffer<f32>> = None;
    let mut coarse: Vec<f32> = Vec::new();
    let mut win_energy = 0f64;
    let mut win_count = 0usize;
    // When the container carries no timestamps, the slice is bounded by
    // counted frames from wherever the seek landed instead.
    let mut frames_seen: u64 = 0;

    loop {
        let packet = match format.next_packet() {
            Ok(p) => p,
            Err(_) => break, // EOF or error → stop with what we have
        };
        if packet.track_id() != track_id {
            continue;
        }
        // Where this packet sits in the file, when the format can say.
        let packet_secs = time_base.map(|tb| {
            let t = tb.calc_time(packet.ts());
            t.seconds as f64 + t.frac
        });
        if let Some(t) = packet_secs {
            if t < start_secs {
                // Before the slice: decode and discard, keeping codec state
                // warm across the boundary.
                let _ = decoder.decode(&packet);
                continue;
            }
            if end_secs.is_some_and(|end| t >= end) {
                break;
            }
        }
        let decoded = match decoder.decode(&packet) {
            Ok(d) => d,
            Err(symphonia::core::errors::Error::DecodeError(_)) => continue,
            Err(_) => break,
        };

        let spec = *decoded.spec();
        let channels = spec.channels.count().max(1);
        if sample_buf.is_none() {
            sample_buf = Some(SampleBuffer::new(decoded.capacity() as u64, spec));
        }
        let sb = sample_buf.as_mut().unwrap();
        sb.copy_interleaved_ref(decoded);

        for frame in sb.samples().chunks(channels) {
            let energy =
                frame.iter().map(|&s| (s as f64) * (s as f64)).sum::<f64>() / channels as f64;
            win_energy += energy;
            win_count += 1;
            if win_count >= WINDOW {
                coarse.push((win_energy / WINDOW as f64).sqrt() as f32);
                win_energy = 0.0;
                win_count = 0;
            }
        }
        if packet_secs.is_none() {
            frames_seen += (sb.samples().len() / channels) as u64;
            if let (Some(end), Some(rate)) = (end_secs, sample_rate) {
                if frames_seen as f64 / rate as f64 >= end - start_secs {
                    break;
                }
            }
        }
    }
    if win_count > 0 {
        coarse.push((win_energy / win_count as f64).sqrt() as f32);
    }

    Ok(coarse)
}

/// Downsample `coarse` RMS values to exactly `buckets`, averaging each
/// span, then normalize so the loudest bucket is 1.0.
fn resample(coarse: &[f32], buckets: usize) -> Vec<f32> {
    let mut out = vec![0f32; buckets];
    if coarse.is_empty() || buckets == 0 {
        return out;
    }
    for (b, slot) in out.iter_mut().enumerate() {
        let start = b * coarse.len() / buckets;
        let end = (((b + 1) * coarse.len() / buckets).max(start + 1)).min(coarse.len());
        let span = &coarse[start..end];
        *slot = span.iter().sum::<f32>() / span.len() as f32;
    }
    let max = out.iter().fold(0f32, |m, &v| m.max(v));
    if max > 0.0 {
        for v in out.iter_mut() {
            *v /= max;
        }
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn resample_normalizes() {
        let coarse = vec![0.2, 0.4, 0.8, 0.1, 0.6, 0.3];
        let out = resample(&coarse, 3);
        assert_eq!(out.len(), 3);
        assert!((out.iter().cloned().fold(0f32, f32::max) - 1.0).abs() < 1e-6);
    }

    #[test]
    fn empty_input() {
        assert_eq!(resample(&[], 5), vec![0.0; 5]);
    }

    /// How much of the track's dynamics survive the trip to the UI?
    ///
    /// The lane looked flat no matter what the frontend did with the numbers,
    /// so this measures where the range actually goes. It decodes a real file
    /// once and re-buckets that single RMS pass four ways, printing the spread
    /// of each and dumping them as JSON for the shaping mockup.
    ///
    /// ```text
    /// $env:SOULSTREAM_WF_FILE = "C:\path\to\track.flac"
    /// $env:SOULSTREAM_WF_DUMP = "C:\path\to\out.json"
    /// cargo test -p player bucketing_strategies -- --ignored --nocapture
    /// ```
    #[test]
    #[ignore = "decodes a real audio file named by SOULSTREAM_WF_FILE"]
    fn bucketing_strategies_on_a_real_track() {
        let Ok(file) = std::env::var("SOULSTREAM_WF_FILE") else {
            panic!("set SOULSTREAM_WF_FILE to an audio file");
        };
        let coarse = coarse_rms(Path::new(&file)).expect("decode");
        println!(
            "{} RMS windows of {WINDOW} frames (~{:.1}s of audio)",
            coarse.len(),
            coarse.len() as f32 * WINDOW as f32 / 44100.0,
        );

        // Each strategy takes the same spans; only the statistic differs.
        fn fold(coarse: &[f32], buckets: usize, pick: impl Fn(&[f32]) -> f32) -> Vec<f32> {
            let mut out = vec![0f32; buckets];
            for (b, slot) in out.iter_mut().enumerate() {
                let start = b * coarse.len() / buckets;
                let end = (((b + 1) * coarse.len() / buckets).max(start + 1)).min(coarse.len());
                *slot = pick(&coarse[start..end]);
            }
            let max = out.iter().fold(0f32, |m, &v| m.max(v));
            if max > 0.0 {
                for v in out.iter_mut() {
                    *v /= max;
                }
            }
            out
        }
        let mean = |s: &[f32]| s.iter().sum::<f32>() / s.len() as f32;
        let peak = |s: &[f32]| s.iter().fold(0f32, |m, &v| m.max(v));
        let blend = move |s: &[f32]| mean(s) * 0.35 + peak(s) * 0.65;

        let mut dump = serde_json::Map::new();
        for (buckets, label) in [(400usize, "400"), (2000usize, "2000")] {
            for (name, values) in [
                ("mean", fold(&coarse, buckets, mean)),
                ("peak", fold(&coarse, buckets, peak)),
                ("blend", fold(&coarse, buckets, blend)),
            ] {
                let mut sorted = values.clone();
                sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
                let q = |f: f32| sorted[((f * (sorted.len() - 1) as f32) as usize).min(sorted.len() - 1)];
                let avg = values.iter().sum::<f32>() / values.len() as f32;
                let sd = (values.iter().map(|v| (v - avg).powi(2)).sum::<f32>()
                    / values.len() as f32)
                    .sqrt();
                println!(
                    "{label:>5} {name:<6} p05 {:.3}  p50 {:.3}  p95 {:.3}  IQR {:.3}  sigma {:.3}",
                    q(0.05),
                    q(0.50),
                    q(0.95),
                    q(0.75) - q(0.25),
                    sd,
                );
                dump.insert(
                    format!("{name}{label}"),
                    serde_json::json!(values.iter().map(|v| (v * 10000.0).round() / 10000.0).collect::<Vec<_>>()),
                );
            }
        }
        if let Ok(out) = std::env::var("SOULSTREAM_WF_DUMP") {
            std::fs::write(&out, serde_json::to_string(&dump).unwrap()).expect("write dump");
            println!("wrote {out}");
        }
    }
}