bass-foundry/plugins/crucible/core/src/take.rs

//! Captured audio with its tempo map, slice boundaries and per-slice cleanup.
use crate::ring::{Marker, Ring, Transport};
use crate::settings::{Division, Mode, Settings};
use common_dsp::{hz_to_midi, Yin};

/// Used when the host reports no tempo.
pub const FALLBACK_TEMPO: f64 = 140.;
/// Waveform overview resolution.
pub const OVERVIEW: usize = 2048;

/// One constant-tempo stretch of the take, relative to the take start.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Segment {
    pub frame: f64,
    pub quarters: f64,
    pub tempo: f64,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Slice {
    pub start: usize,
    pub end: usize,
    pub enabled: bool,
    pub root: Option<u8>,
}

#[derive(Clone, Debug)]
pub struct Take {
    pub rate: f32,
    /// Interleaved stereo.
    pub audio: Vec<f32>,
    pub bars: u32,
    pub numerator: u32,
    pub denominator: u32,
    /// True when aligned to the host's bar lines; false when the transport was stopped.
    pub synced: bool,
    pub segments: Vec<Segment>,
    /// Slice boundaries in frames, sorted, starting at 0 and ending at the take length.
    pub cuts: Vec<usize>,
    pub slices: Vec<Slice>,
    /// Min/max per overview bucket of the mono sum.
    pub overview: Vec<[f32; 2]>,
}

impl Take {
    pub fn frames(&self) -> usize {
        self.audio.len() / 2
    }
    pub fn tempo(&self) -> f64 {
        self.segments.first().map_or(FALLBACK_TEMPO, |s| s.tempo)
    }
    pub fn bar_quarters(&self) -> f64 {
        self.numerator.max(1) as f64 * 4. / self.denominator.max(1) as f64
    }
    pub fn quarters(&self) -> f64 {
        self.bars as f64 * self.bar_quarters()
    }
    /// Frame at a position in quarter notes from the take start.
    pub fn frame_at(&self, quarters: f64) -> f64 {
        let s = self
            .segments
            .iter()
            .rev()
            .find(|s| s.quarters <= quarters + 1e-9)
            .or(self.segments.first());
        match s {
            Some(s) => s.frame + (quarters - s.quarters) * 60. / s.tempo * self.rate as f64,
            None => quarters * 60. / FALLBACK_TEMPO * self.rate as f64,
        }
    }
    pub fn mono(&self, start: usize, end: usize) -> Vec<f32> {
        self.audio[start * 2..end * 2]
            .as_chunks::<2>()
            .0
            .iter()
            .map(|f| 0.5 * (f[0] + f[1]))
            .collect()
    }
    fn build_overview(&mut self) {
        let n = self.frames();
        self.overview = (0..OVERVIEW)
            .map(|b| {
                let (a, z) = (
                    b * n / OVERVIEW,
                    ((b + 1) * n / OVERVIEW).max(b * n / OVERVIEW + 1),
                );
                let mut mm = [0f32, 0f32];
                for f in a..z.min(n) {
                    let v = 0.5 * (self.audio[2 * f] + self.audio[2 * f + 1]);
                    mm[0] = mm[0].min(v);
                    mm[1] = mm[1].max(v);
                }
                mm
            })
            .collect();
    }
    pub fn from_audio(
        rate: f32,
        audio: Vec<f32>,
        bars: u32,
        transport: Transport,
        segments: Vec<Segment>,
        synced: bool,
    ) -> Self {
        let mut t = Take {
            rate,
            audio,
            bars,
            numerator: transport.numerator.max(1),
            denominator: transport.denominator.max(1),
            synced,
            segments,
            cuts: Vec::new(),
            slices: Vec::new(),
            overview: Vec::new(),
        };
        t.build_overview();
        let n = t.frames();
        t.cuts = vec![0, n];
        t.rebuild_slices();
        t
    }
    /// Re-cuts the take from the settings, replacing manual edits.
    pub fn slice(&mut self, settings: &Settings) {
        let n = self.frames();
        let mut cuts = match settings.mode {
            Mode::Grid => grid(self, settings.division),
            Mode::Transient => transients(&self.mono(0, n), self.rate, settings.sensitivity),
        };
        cuts.push(0);
        cuts.push(n);
        cuts.sort_unstable();
        cuts.dedup();
        cuts.retain(|c| *c <= n);
        self.cuts = cuts;
        self.rebuild_slices();
    }
    /// Adds a cut; returns false if it is too close to an existing one.
    pub fn add_cut(&mut self, frame: usize) -> bool {
        let min = (self.rate * 0.005) as usize;
        if frame == 0 || frame >= self.frames() || self.cuts.iter().any(|c| c.abs_diff(frame) < min)
        {
            return false;
        }
        let at = self.cuts.partition_point(|c| *c < frame);
        self.cuts.insert(at, frame);
        self.rebuild_slices();
        true
    }
    /// Removes the inner cut nearest to `frame`.
    pub fn remove_cut(&mut self, frame: usize) -> bool {
        let inner = 1..self.cuts.len().saturating_sub(1);
        let Some(i) = inner.min_by_key(|i| self.cuts[*i].abs_diff(frame)) else {
            return false;
        };
        self.cuts.remove(i);
        self.rebuild_slices();
        true
    }
    fn rebuild_slices(&mut self) {
        let old = std::mem::take(&mut self.slices);
        let mut yin = Yin::new((self.rate / 35.) as usize + 4);
        self.slices = self
            .cuts
            .windows(2)
            .map(|w| {
                // Manual include/exclude choices survive a re-cut when the slice keeps its start.
                let enabled = old
                    .iter()
                    .find(|s| s.start == w[0])
                    .is_none_or(|s| s.enabled);
                let mono = self.mono(w[0], w[1]);
                Slice {
                    start: w[0],
                    end: w[1],
                    enabled,
                    root: root_key(&mut yin, &mono, self.rate),
                }
            })
            .collect();
    }
    pub fn slice_at(&self, frame: usize) -> Option<usize> {
        self.slices
            .iter()
            .position(|s| frame >= s.start && frame < s.end)
    }
}

fn grid(take: &Take, division: Division) -> Vec<usize> {
    let per = division.per_quarter() as f64;
    let steps = (take.quarters() * per).round() as usize;
    (1..steps)
        .map(|k| take.frame_at(k as f64 / per).round() as usize)
        .collect()
}

/// Onsets where a window's high-frequency weighted energy jumps above the window before it.
/// Windows span 40 ms, longer than one cycle of the lowest bass note, so the edges inside a
/// distorted bass cycle never count as hits. Each onset is refined to its first loud sample.
pub fn transients(x: &[f32], rate: f32, sensitivity: f32) -> Vec<usize> {
    let hop = ((rate / 375.) as usize).max(32);
    let window = ((rate * 0.04) as usize / hop).max(1) * hop;
    if x.len() < 2 * window + hop {
        return Vec::new();
    }
    // Prefix sums of weighted energy, in f64 so long takes keep their precision.
    let mut sums = Vec::with_capacity(x.len() + 1);
    sums.push(0f64);
    let mut prev = 0f32;
    for v in x {
        let d = v - prev;
        prev = *v;
        let e = (v * v + 4. * d * d) as f64;
        sums.push(sums.last().unwrap() + e);
    }
    let level = |end: usize| {
        let e = (sums[end] - sums[end - window]) / window as f64;
        10. * (e + 1e-10).log10() as f32
    };
    let rise_db = 14. - 0.11 * sensitivity.clamp(0., 100.);
    let floor = -55.;
    let gap = (0.04 * rate) as usize;
    let mut cuts: Vec<usize> = Vec::new();
    let mut end = 2 * window;
    while end <= x.len() {
        let now = level(end);
        if now > floor && now - level(end - window) >= rise_db {
            let from = end - window;
            let to = (end + hop).min(x.len());
            let peak = x[from..to].iter().fold(0f32, |m, v| m.max(v.abs()));
            let first = x[from..to]
                .iter()
                .position(|v| v.abs() >= 0.3 * peak)
                .unwrap_or(0);
            // Start slightly before the first loud sample so the attack stays whole.
            let onset = (from + first).saturating_sub((rate * 0.0005) as usize);
            if cuts.last().is_none_or(|c| onset >= c + gap) {
                cuts.push(onset);
                // Skip past this attack so it is not found again as the window slides on.
                end = onset + window + gap;
            }
        }
        end += hop;
    }
    cuts
}

/// Root note of a slice, if it is clearly pitched.
pub fn root_key(yin: &mut Yin, mono: &[f32], rate: f32) -> Option<u8> {
    let window = (2 * yin.max_lag() + 2).min(mono.len());
    if window < (rate / 1000.) as usize * 4 || mono.len() < window {
        return None;
    }
    // Skip the attack: the body of a hit carries its pitch.
    let mut notes: Vec<f32> = Vec::new();
    let start = (mono.len() / 5).min(mono.len() - window);
    let step = (window / 2).max(1);
    let mut at = start;
    while at + window <= mono.len() && notes.len() < 6 {
        if let Some((hz, clarity)) = yin.estimate(&mono[at..at + window], rate, 35., 1000.) {
            if clarity > 0.7 {
                notes.push(hz_to_midi(hz));
            }
        }
        at += step;
    }
    if notes.is_empty() {
        return None;
    }
    notes.sort_by(f32::total_cmp);
    let n = notes[notes.len() / 2].round();
    (0. ..=127.).contains(&n).then_some(n as u8)
}

/// Builds a take from the newest `bars` complete bars, or from the newest audio when stopped.
pub fn capture(ring: &Ring, bars: u32) -> Result<Take, &'static str> {
    let rate = ring.rate() as f64;
    let written = ring.written();
    let markers = ring.markers();
    let last = markers.iter().rev().find(|m| m.frame <= written).copied();
    let fallback = Transport {
        playing: false,
        tempo: FALLBACK_TEMPO,
        beats: 0.,
        numerator: 4,
        denominator: 4,
    };
    let t = last.map_or(fallback, |m| m.transport);
    let bar = t.bar();
    let quarters = bars as f64 * bar;
    let (start, end, segments, synced) = match last {
        Some(m) if m.transport.playing => {
            let now = m.beats_at(written as f64, rate);
            let end_beats = (now / bar + 1e-9).floor() * bar;
            let start_beats = end_beats - quarters;
            let end = frame_for(&markers, end_beats, rate);
            let start = frame_for(&markers, start_beats, rate);
            let segments = segments(&markers, start_beats, end_beats, start, rate);
            (start, end, segments, true)
        }
        _ => {
            let length = quarters * 60. / t.tempo * rate;
            let segments = vec![Segment {
                frame: 0.,
                quarters: 0.,
                tempo: t.tempo,
            }];
            (written as f64 - length, written as f64, segments, false)
        }
    };
    if start < 0. {
        return Err("Not enough audio recorded yet. Play the section first.");
    }
    let (start, end) = (start.round() as u64, end.round() as u64);
    if end - start > ring.capacity() as u64 {
        return Err("Too long for the recording buffer. Capture fewer bars.");
    }
    let audio = ring.copy(start, end)?;
    Ok(Take::from_audio(
        ring.rate(),
        audio,
        bars,
        t,
        segments,
        synced,
    ))
}

/// Frame of a song position, from the marker whose stretch contains it.
fn frame_for(markers: &[Marker], beats: f64, rate: f64) -> f64 {
    let m = markers
        .iter()
        .rev()
        .find(|m| m.transport.playing && m.transport.beats <= beats + 1e-9)
        .or_else(|| markers.iter().find(|m| m.transport.playing));
    m.map_or(0., |m| m.frame_at(beats, rate))
}

fn segments(markers: &[Marker], from: f64, to: f64, origin: f64, rate: f64) -> Vec<Segment> {
    let mut out: Vec<Segment> = Vec::new();
    for m in markers.iter().filter(|m| m.transport.playing) {
        let q = m.transport.beats - from;
        if m.transport.beats >= to {
            break;
        }
        let tempo = m.transport.tempo;
        if q <= 0. {
            out.clear();
            out.push(Segment {
                frame: m.frame_at(from, rate) - origin,
                quarters: 0.,
                tempo,
            });
        } else if out.last().is_none_or(|s| (s.tempo - tempo).abs() > 1e-6) {
            out.push(Segment {
                frame: m.frame as f64 - origin,
                quarters: q,
                tempo,
            });
        }
    }
    if out.is_empty() {
        out.push(Segment {
            frame: 0.,
            quarters: 0.,
            tempo: FALLBACK_TEMPO,
        });
    }
    out
}

/// A cleaned-up slice ready to save: trimmed, faded and optionally normalised.
pub fn render(take: &Take, slice: &Slice, settings: &Settings) -> Vec<f32> {
    let mut start = slice.start;
    let end = slice.end;
    if settings.trim {
        let threshold = 10f32.powf(-60. / 20.);
        if let Some(i) = take.audio[start * 2..end * 2]
            .as_chunks::<2>()
            .0
            .iter()
            .position(|f| f[0].abs().max(f[1].abs()) > threshold)
        {
            // Keep a short pre-roll so the attack is not cut into.
            start += i.saturating_sub((take.rate * 0.001) as usize);
        }
    }
    let mut out = take.audio[start * 2..end * 2].to_vec();
    let frames = out.len() / 2;
    let fade = ((settings.declick * 0.001 * take.rate) as usize).min(frames / 2);
    for i in 0..fade {
        let g = i as f32 / fade as f32;
        for c in 0..2 {
            out[2 * i + c] *= g;
            out[2 * (frames - 1 - i) + c] *= g;
        }
    }
    if settings.normalize {
        let peak = out.iter().fold(0f32, |m, v| m.max(v.abs()));
        if peak > 1e-6 {
            let g = 10f32.powf(-0.3 / 20.) / peak;
            out.iter_mut().for_each(|v| *v *= g);
        }
    }
    out
}

pub const TABLE: usize = 2048;
pub const TABLE_FRAMES: usize = 256;

/// Single-cycle frames from a pitched slice. Each frame starts on a rising zero crossing,
/// so neighbouring frames stay in phase and morph smoothly in a wavetable synth.
pub fn wavetable(mono: &[f32], rate: f32) -> Option<Vec<f32>> {
    let mut yin = Yin::new((rate / 30.) as usize + 4);
    let window = 2 * yin.max_lag() + 2;
    let probe = &mono[(mono.len() / 5).min(mono.len().saturating_sub(window))..];
    let (hz, _) = yin.estimate(&probe[..window.min(probe.len())], rate, 30., 2000.)?;
    let period = rate as f64 / hz as f64;
    let rising = |i: usize| mono[i - 1] < 0. && mono[i] >= 0.;
    let crossing = |i: usize| (i - 1) as f64 + mono[i - 1] as f64 / (mono[i - 1] - mono[i]) as f64;
    let mut starts = Vec::new();
    let mut at = (1..mono.len()).find(|i| rising(*i))?;
    starts.push(crossing(at));
    loop {
        let expect = starts.last().unwrap() + period;
        let tolerance = period * 0.2;
        let lo = (expect - tolerance).max(1.) as usize;
        let hi = ((expect + tolerance) as usize).min(mono.len() - 1);
        if lo >= hi {
            break;
        }
        // The rising crossing nearest one period on; slight drift and detune are followed.
        let Some(next) = (lo..=hi).filter(|i| rising(*i)).min_by(|a, b| {
            (crossing(*a) - expect)
                .abs()
                .total_cmp(&(crossing(*b) - expect).abs())
        }) else {
            // No crossing: skip one period and look again.
            at = (expect as usize).min(mono.len() - 1);
            match (at..mono.len()).find(|i| rising(*i)) {
                Some(i) => {
                    starts.push(crossing(i));
                    continue;
                }
                None => break,
            }
        };
        starts.push(crossing(next));
    }
    let cycles: Vec<(f64, f64)> = starts
        .windows(2)
        .map(|w| (w[0], w[1]))
        .filter(|(a, b)| (b - a - period).abs() < period * 0.25)
        .collect();
    if cycles.is_empty() {
        return None;
    }
    let count = cycles.len().min(TABLE_FRAMES);
    let mut table = Vec::with_capacity(count * TABLE);
    for k in 0..count {
        let (a, b) = cycles[k * cycles.len() / count];
        for j in 0..TABLE {
            let pos = a + (b - a) * j as f64 / TABLE as f64;
            let i = pos.floor() as usize;
            let frac = (pos - i as f64) as f32;
            let x0 = mono[i.min(mono.len() - 1)];
            let x1 = mono[(i + 1).min(mono.len() - 1)];
            table.push(x0 + (x1 - x0) * frac);
        }
    }
    let peak = table.iter().fold(0f32, |m, v| m.max(v.abs()));
    if peak < 1e-6 {
        return None;
    }
    let g = 10f32.powf(-0.3 / 20.) / peak;
    table.iter_mut().for_each(|v| *v *= g);
    Some(table)
}