OnsetDetector.h

// DSPark -- Professional Audio DSP Framework
// Copyright (c) 2026 Cristian Moresi -- MIT License
 
#pragma once
 
#include "../Core/DspMath.h"
#include "detail/OnsetFeatures.h"
#include "../Core/AudioBuffer.h"
#include "../Core/AudioSpec.h"
 
#include <algorithm>
#include <array>
#include <atomic>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <span>
#include <vector>
 
namespace dspark {
 
template <FloatType T>
class OnsetDetector final
{
public:
    enum class Method { SpectralFlux, ComplexDomain, SuperFlux };
 
    // -- Lifecycle -----------------------------------------------------------
 
    void prepare(const AudioSpec& spec, int fftSize = 0, int hop = 0)
    {
        if (!(spec.sampleRate > 0.0) || !std::isfinite(spec.sampleRate))
            return;
 
        prepared_ = false; // gate OFF: the audio path is a no-op while rebuilding
 
        sampleRate_ = spec.sampleRate;
 
        if (fftSize <= 0)
        {
            // Automatic: hold the analysis TIME SPAN constant across sample
            // rates (the span 2048 samples cover at 48 kHz). Bin width and the
            // filterbank's quarter-tone floor are both fs/fftSize, so a
            // constant span pins them -- and with them the usable low register
            // -- instead of letting them widen as fs rises.
            fftSize_ = detail::OnsetSpectrum<T>::defaultFrameSize(sampleRate_);
        }
        else
        {
            int fs = std::clamp(fftSize, kMinFft, kMaxFft);
            int pow2 = kMinFft;
            while (pow2 < fs) pow2 <<= 1;
            fftSize_ = pow2;
        }
 
        if (hop <= 0)
            hop_ = std::max(1, static_cast<int>(std::lround(sampleRate_ / 200.0)));
        else
            hop_ = hop;
        hop_ = std::clamp(hop_, 1, fftSize_);
 
        latencySamples_.store(static_cast<int64_t>(fftSize_) + hop_,
                              std::memory_order_relaxed);
 
        // Analysis-window group-delay compensation: half-wave-rectified spectral
        // flux peaks on the rising flank of a transient, ~kLocalizationLead*N
        // before the windowed-energy peak. Adding it back centres the reported
        // onset on the transient (calibrated on band-limited clicks; softer
        // onsets localise slightly late but stay within the acceptance window).
        localizationOffset_ = static_cast<int>(std::lround(kLocalizationLead
                                                           * static_cast<double>(fftSize_)));
        // Warm-up: suppress onsets until the analysis ring is fully primed so
        // the silence->first-input ramp cannot fire a spurious onset.
        primeFrames_ = fftSize_ / hop_ + 2;
 
        ring_.assign(static_cast<size_t>(fftSize_) * 2, T(0));
        spectrum_.prepare(fftSize_);
        novelty_.prepare(sampleRate_, fftSize_);
 
        // Peak-picker windows in frames (derived from the ms defaults above).
        preMaxFrames_  = msToFrames(30.0);
        postMaxFrames_ = msToFrames(30.0);   // offline; causal uses 1
        preAvgFrames_  = msToFrames(100.0);
        postAvgFrames_ = msToFrames(70.0);   // offline; causal uses 0
        waitFrames_    = msToFrames(30.0);
 
        // Causal ODF history: enough for the pre_avg look-back plus the
        // single-frame causal confirmation.
        odfHistLen_ = std::max(preAvgFrames_, preMaxFrames_) + 4;
        odfHist_.assign(static_cast<size_t>(odfHistLen_), T(0));
 
        // Pending-onset ring: onsets are held from detection until their
        // release sample (reference + L). At most ceil(L/hop)+2 can be
        // in flight; size generously to a power-of-two-ish bound.
        const int maxPending = (fftSize_ + hop_) / hop_ + 4;
        pending_.assign(static_cast<size_t>(std::max(8, maxPending)),
                        PendingOnset{});
        pendingHead_ = 0;
        pendingCount_ = 0;
 
        resetState();
 
        method_.store(Method::SuperFlux, std::memory_order_relaxed);
        // A conservative default delta, tuned against a synthetic corpus of
        // clicks, soft onsets and noise beds. Callers override per material.
        threshold_.store(kDefaultDelta, std::memory_order_relaxed);
        whitening_.store(false, std::memory_order_relaxed);
 
        prepared_ = true; // gate ON
    }
 
    void setMethod(Method m) noexcept { method_.store(m, std::memory_order_relaxed); }
 
    void setThreshold(T deltaAboveMean) noexcept
    {
        if (!std::isfinite(deltaAboveMean)) return;
        threshold_.store(std::max(T(0), deltaAboveMean), std::memory_order_relaxed);
    }
 
    void setAdaptiveWhitening(bool on) noexcept
    {
        whitening_.store(on, std::memory_order_relaxed);
    }
 
    [[nodiscard]] Method getMethod() const noexcept
    {
        return method_.load(std::memory_order_relaxed);
    }
 
    [[nodiscard]] T getThreshold() const noexcept
    {
        return threshold_.load(std::memory_order_relaxed);
    }
 
    [[nodiscard]] bool getAdaptiveWhitening() const noexcept
    {
        return whitening_.load(std::memory_order_relaxed);
    }
 
    // -- Audio path (causal, RT-safe) ---------------------------------------
 
    void processBlock(AudioBufferView<const T> in) noexcept
    {
        if (!prepared_ || in.getNumChannels() < 1) return;
        const T* ch0 = in.getChannel(0);
        pushSamples(std::span<const T>(ch0, static_cast<size_t>(in.getNumSamples())));
    }
 
    void pushSamples(std::span<const T> samples) noexcept
    {
        if (!prepared_) return;
 
        bool firedThisCall = false;
 
        for (const T s : samples)
        {
            const T x = std::isfinite(s) ? s : T(0);
 
            ring_[static_cast<size_t>(writePos_)] = x;
            ring_[static_cast<size_t>(writePos_ + fftSize_)] = x;
            if (++writePos_ >= fftSize_) writePos_ = 0;
 
            ++totalSamples_;
 
            // Release any pending onset whose reporting sample has arrived.
            while (pendingCount_ > 0)
            {
                const PendingOnset& p = pending_[static_cast<size_t>(pendingHead_)];
                if (totalSamples_ - 1 >= p.reportSample)
                {
                    lastOnsetSample_.store(p.referenceSample, std::memory_order_relaxed);
                    onsetStrength_.store(p.strength, std::memory_order_relaxed);
                    firedThisCall = true;
                    pendingHead_ = (pendingHead_ + 1) % static_cast<int>(pending_.size());
                    --pendingCount_;
                }
                else break;
            }
 
            if (++hopCounter_ >= hop_)
            {
                hopCounter_ = 0;
                analyzeFrame();
            }
        }
 
        onsetLatched_.store(firedThisCall, std::memory_order_relaxed);
    }
 
    // -- Readout (lock-free) -------------------------------------------------
 
    [[nodiscard]] bool onsetDetected() const noexcept
    {
        return onsetLatched_.load(std::memory_order_relaxed);
    }
 
    [[nodiscard]] T getOnsetStrength() const noexcept
    {
        return onsetStrength_.load(std::memory_order_relaxed);
    }
 
    [[nodiscard]] int64_t getLastOnsetSample() const noexcept
    {
        return lastOnsetSample_.load(std::memory_order_relaxed);
    }
 
    [[nodiscard]] int getLatencySamples() const noexcept
    {
        return static_cast<int>(latencySamples_.load(std::memory_order_relaxed));
    }
 
    [[nodiscard]] int getFftSize() const noexcept { return fftSize_; }
 
    [[nodiscard]] int getHopSize() const noexcept { return hop_; }
 
    [[nodiscard]] int getNumBands() const noexcept { return novelty_.getNumBands(); }
 
    // -- Onset-strength envelope (stream owner only) -------------------------
 
    static constexpr int kNumRegisters = detail::OnsetNovelty<T>::kNumRegisters;
 
    struct OdfFrame
    {
        T value = T(0);              
        int64_t referenceSample = 0; 
        int64_t frameIndex = 0;      
        std::array<T, kNumRegisters> registers {};
    };
 
    [[nodiscard]] OdfFrame getLastOdfFrame() const noexcept
    {
        return OdfFrame { lastOdfValue_, lastOdfRef_, frameIndex_, lastRegisters_ };
    }
 
    [[nodiscard]] int getWarmupFrames() const noexcept { return primeFrames_; }
 
    [[nodiscard]] int getEnvelopeLatencySamples() const noexcept
    {
        return fftSize_ / 2 - localizationOffset_;
    }
 
    // -- Offline -------------------------------------------------------------
 
    struct Onset
    {
        int64_t sample = 0; 
        T strength = T(0);  
    };
 
    std::vector<int64_t> detectOffline(AudioBufferView<const T> whole)
    {
        return positionsOf(detectOfflineOnsets(whole));
    }
 
    std::vector<Onset> detectOfflineOnsets(AudioBufferView<const T> whole)
    {
        if (!prepared_ || whole.getNumChannels() < 1) return {};
        const int n = std::max(0, whole.getNumSamples());
        beginOffline(n);
        pushOffline(std::span<const T>(whole.getChannel(0), static_cast<size_t>(n)));
        return finishOfflineOnsets();
    }
 
    void beginOffline(int64_t expectedSamples = 0)
    {
        offOdf_.clear();
        offRef_.clear();
        offlineOpen_ = false;
        if (!prepared_) return;
        if (expectedSamples > 0)
        {
            const size_t frames = static_cast<size_t>(expectedSamples / std::max(1, hop_) + 2);
            offOdf_.reserve(frames);
            offRef_.reserve(frames);
        }
        resetState();
        offMethod_ = method_.load(std::memory_order_relaxed);
        offWhiten_ = whitening_.load(std::memory_order_relaxed);
        offTotal_ = 0;
        offHop_ = 0;
        offlineOpen_ = true;
    }
 
    void pushOffline(std::span<const T> samples)
    {
        if (!offlineOpen_) return;
        for (const T s : samples)
        {
            const T v = std::isfinite(s) ? s : T(0);
            ring_[static_cast<size_t>(writePos_)] = v;
            ring_[static_cast<size_t>(writePos_ + fftSize_)] = v;
            if (++writePos_ >= fftSize_) writePos_ = 0;
            ++offTotal_;
            if (++offHop_ >= hop_)
            {
                offHop_ = 0;
                offOdf_.push_back(computeOdf(offMethod_, offWhiten_));
                offRef_.push_back(referenceSample(offTotal_));
            }
        }
    }
 
    std::vector<int64_t> finishOffline()
    {
        return positionsOf(finishOfflineOnsets());
    }
 
    std::vector<Onset> finishOfflineOnsets()
    {
        std::vector<Onset> out;
        if (!offlineOpen_) return out;
        offlineOpen_ = false;
 
        // Symmetric peak-pick over the whole envelope.
        const T delta = threshold_.load(std::memory_order_relaxed);
        const std::vector<T>& odf = offOdf_;
        detail::OnsetPeakPicker<T> picker;
        const typename detail::OnsetPeakPicker<T>::Windows windows {
            static_cast<size_t>(preMaxFrames_), static_cast<size_t>(postMaxFrames_),
            static_cast<size_t>(preAvgFrames_), static_cast<size_t>(postAvgFrames_),
            static_cast<size_t>(waitFrames_)
        };
        for (size_t f = static_cast<size_t>(primeFrames_); f < odf.size(); ++f)
        {
            if (picker.accept(f, odf.size(), [&](size_t i) { return odf[i]; }, windows, delta))
                out.push_back(Onset { offRef_[f], odf[f] });
        }
 
        offOdf_.clear();
        offRef_.clear();
        resetState();
        return out;
    }
 
    void reset() noexcept
    {
        offlineOpen_ = false;
        resetState();
    }
 
private:
    [[nodiscard]] static std::vector<int64_t> positionsOf(const std::vector<Onset>& onsets)
    {
        std::vector<int64_t> out;
        out.reserve(onsets.size());
        for (const Onset& o : onsets) out.push_back(o.sample);
        return out;
    }
 
    // -- Constants -----------------------------------------------------------
    static constexpr int kMinFft = 64;
    static constexpr int kMaxFft = 1 << 16;
    static constexpr int64_t kBig = int64_t(1) << 60;
    static constexpr T kDefaultDelta = T(0.03);
    static constexpr double kLocalizationLead = 0.34; 
 
    struct PendingOnset
    {
        int64_t referenceSample = 0; 
        int64_t reportSample = 0;    
        T strength = T(0);
    };
 
    // -- Frame timing --------------------------------------------------------
 
    [[nodiscard]] int64_t referenceSample(int64_t firePos) const noexcept
    {
        return firePos - static_cast<int64_t>(fftSize_) / 2
             + static_cast<int64_t>(localizationOffset_);
    }
 
    [[nodiscard]] int msToFrames(double ms) const noexcept
    {
        return std::max(1, static_cast<int>(std::lround(ms * 0.001 * sampleRate_
                                                        / static_cast<double>(hop_))));
    }
 
    // -- STFT + ODF ----------------------------------------------------------
 
    T computeOdf(Method method, bool whiten) noexcept
    {
        spectrum_.compute(&ring_[static_cast<size_t>(writePos_)]);
        using FeatureMethod = typename detail::OnsetNovelty<T>::Method;
        const auto selected = method == Method::SpectralFlux ? FeatureMethod::SpectralFlux
                            : method == Method::ComplexDomain ? FeatureMethod::ComplexDomain
                            : method == Method::SuperFlux ? FeatureMethod::SuperFlux
                            : static_cast<FeatureMethod>(-1);
        const auto result = novelty_.process(spectrum_.magnitudes(), spectrum_.phases(),
            selected, whiten);
        curRegisters_ = result.registers;
        return result.value;
    }
 
    void analyzeFrame() noexcept
    {
        const Method m = method_.load(std::memory_order_relaxed);
        const bool whiten = whitening_.load(std::memory_order_relaxed);
        const T value = computeOdf(m, whiten);
 
        // Push into the causal ODF history ring.
        odfHist_[static_cast<size_t>(odfWrite_)] = value;
        odfWrite_ = (odfWrite_ + 1) % odfHistLen_;
        ++frameIndex_;
 
        // Envelope readout for the stream owner (see getLastOdfFrame()). The
        // frame fired at totalSamples_, so it localises exactly where an onset
        // decided from it would: one definition of frame time, not two.
        lastOdfValue_ = value;
        lastOdfRef_ = referenceSample(totalSamples_);
        lastRegisters_ = curRegisters_;
 
        // Causal peak-pick with a single-frame confirmation: we decide whether
        // the PREVIOUS frame was a maximum now that we have the current one.
        // This one-hop confirmation is exactly the +hop term of L.
        if (frameIndex_ < 2) { lastConfirmOdf_ = value; return; }
 
        const T prev = odfAt(1);      // previous frame's ODF (candidate)
        const T curr = value;         // current frame (confirmation)
 
        // (1) prev is a causal local max over [prev-preMax, prev+1].
        bool isMax = (prev >= curr);
        if (isMax)
        {
            for (int j = 2; j <= preMaxFrames_ + 1 && j < frameIndex_; ++j)
            {
                if (odfAt(j) > prev) { isMax = false; break; }
            }
        }
 
        if (isMax)
        {
            // (2) prev exceeds the backward moving mean + delta.
            T sum = T(0); int cnt = 0;
            for (int j = 1; j <= preAvgFrames_ && j < frameIndex_; ++j)
            {
                sum += odfAt(j); ++cnt;
            }
            const T mean = (cnt > 0) ? sum / static_cast<T>(cnt) : T(0);
            const T delta = threshold_.load(std::memory_order_relaxed);
 
            // (3) combination width since the last onset, and past warm-up.
            const int64_t candFrame = frameIndex_ - 1;
            if (candFrame >= primeFrames_ && prev >= mean + delta
                && candFrame - lastOnsetFrame_ > waitFrames_)
            {
                // The candidate is the previous frame; its fire position was
                // totalSamples_ - hop (one hop before the current frame's fire).
                const int64_t candFirePos = totalSamples_ - hop_;
                const int64_t ref = referenceSample(candFirePos);
                scheduleOnset(ref, prev);
                lastOnsetFrame_ = candFrame;
            }
        }
    }
 
    [[nodiscard]] T odfAt(int j) const noexcept
    {
        int idx = odfWrite_ - 1 - j;
        idx %= odfHistLen_;
        if (idx < 0) idx += odfHistLen_;
        return odfHist_[static_cast<size_t>(idx)];
    }
 
    void scheduleOnset(int64_t referenceSample, T strength) noexcept
    {
        if (pendingCount_ >= static_cast<int>(pending_.size())) return; // saturate
        const int tail = (pendingHead_ + pendingCount_) % static_cast<int>(pending_.size());
        PendingOnset& p = pending_[static_cast<size_t>(tail)];
        p.referenceSample = referenceSample;
        p.reportSample = referenceSample + latencySamples_.load(std::memory_order_relaxed);
        p.strength = strength;
        ++pendingCount_;
    }
 
    void resetState() noexcept
    {
        std::fill(ring_.begin(), ring_.end(), T(0));
        novelty_.reset();
        std::fill(odfHist_.begin(), odfHist_.end(), T(0));
 
        writePos_ = 0;
        hopCounter_ = 0;
        totalSamples_ = 0;
        frameIndex_ = 0;
        odfWrite_ = 0;
        lastOdfValue_ = T(0);
        lastOdfRef_ = 0;
        lastRegisters_.fill(T(0));
        curRegisters_.fill(T(0));
        lastConfirmOdf_ = T(0);
        lastOnsetFrame_ = -kBig;
        pendingHead_ = 0;
        pendingCount_ = 0;
 
        onsetLatched_.store(false, std::memory_order_relaxed);
        onsetStrength_.store(T(0), std::memory_order_relaxed);
        lastOnsetSample_.store(-1, std::memory_order_relaxed);
    }
 
    // -- Members -------------------------------------------------------------
    double sampleRate_ = 44100.0;
    int fftSize_ = 2048;
    int hop_ = 221;
    int localizationOffset_ = 0;
    int primeFrames_ = 12;
    bool prepared_ = false;
    detail::OnsetSpectrum<T> spectrum_;
    detail::OnsetNovelty<T> novelty_;
    std::vector<T> ring_; // Mirrored chronological input, size 2*fftSize.
 
    // Peak-picker windows (frames).
    int preMaxFrames_ = 6, postMaxFrames_ = 6;
    int preAvgFrames_ = 20, postAvgFrames_ = 14;
    int waitFrames_ = 6;
 
    // Causal ODF history ring.
    std::vector<T> odfHist_;
    int odfHistLen_ = 32;
    int odfWrite_ = 0;
    int64_t frameIndex_ = 0;
    T lastOdfValue_ = T(0);   
    int64_t lastOdfRef_ = 0;  
    std::array<T, kNumRegisters> lastRegisters_ {};  
    std::array<T, kNumRegisters> curRegisters_ {};   
    T lastConfirmOdf_ = T(0);
    int64_t lastOnsetFrame_ = -kBig;
 
    // Streaming counters.
    int writePos_ = 0;
    int hopCounter_ = 0;
    int64_t totalSamples_ = 0;
 
    // Pending-onset ring (held from detection to report point).
    std::vector<PendingOnset> pending_;
    int pendingHead_ = 0;
    int pendingCount_ = 0;
 
    // Atomic parameters / readouts.
    std::atomic<Method> method_ { Method::SuperFlux };
    std::atomic<T> threshold_ { kDefaultDelta };
    std::atomic<bool> whitening_ { false };
    std::atomic<bool> onsetLatched_ { false };
    std::atomic<T> onsetStrength_ { T(0) };
    std::atomic<int64_t> lastOnsetSample_ { -1 };
    std::atomic<int64_t> latencySamples_ { 2269 };
 
    // Offline session (beginOffline() .. finishOffline()).
    std::vector<T> offOdf_;
    std::vector<int64_t> offRef_;
    int64_t offTotal_ = 0;
    int offHop_ = 0;
    Method offMethod_ = Method::SuperFlux;
    bool offWhiten_ = false;
    bool offlineOpen_ = false;
};
 
} // namespace dspark