DelayEstimator.h

// DSPark - Professional Audio DSP Framework
// Copyright (c) 2026 Cristian Moresi - MIT License
 
#pragma once
 
#include "../Core/DspMath.h"
#include "../Core/AudioBuffer.h"
#include "../Core/FFT.h"
 
#include <algorithm>
#include <cmath>
#include <complex>
#include <cstdint>
#include <memory>
#include <vector>
 
namespace dspark {
 
template <FloatType T>
class DelayEstimator
{
public:
    struct Result
    {
        double delaySamples = 0.0;  
        double delaySeconds = 0.0;  
        double confidence = 0.0;    
        bool inverted = false;      
        bool valid = false;         
    };
 
    void prepare(double sampleRate, int maxDelaySamples)
    {
        if (!(sampleRate > 0.0) || !std::isfinite(sampleRate)) return;
        sampleRate_ = sampleRate;
        maxDelay_ = std::max(1, maxDelaySamples);
        frame_ = 4096;
        while (frame_ < 8 * maxDelay_ && frame_ < (1 << 26)) frame_ <<= 1;
        hop_ = frame_ / 2;
        fft_ = std::make_unique<FFTReal<double>>(static_cast<size_t>(frame_));
        window_.resize(static_cast<size_t>(frame_));
        for (int i = 0; i < frame_; ++i)
            window_[static_cast<size_t>(i)] = 0.5 - 0.5 * std::cos(2.0 * pi<double> * i / frame_);
        // The reference needs one frame; the delivery a frame plus the delay
        // it is framed at, either side.
        historyLength_ = frame_ + 2 * maxDelay_ + 2;
        refHistory_.assign(static_cast<size_t>(historyLength_), 0.0);
        delHistory_.assign(static_cast<size_t>(historyLength_), 0.0);
        bufA_.assign(static_cast<size_t>(frame_ + 2), 0.0);
        bufB_.assign(static_cast<size_t>(frame_ + 2), 0.0);
        cross_.assign(static_cast<size_t>(frame_ / 2 + 1), std::complex<double> {});
        corr_.assign(static_cast<size_t>(frame_ + 2), 0.0);
        prepared_ = true;
        reset();
    }
 
    void reset() noexcept
    {
        std::fill(refHistory_.begin(), refHistory_.end(), 0.0);
        std::fill(delHistory_.begin(), delHistory_.end(), 0.0);
        std::fill(cross_.begin(), cross_.end(), std::complex<double> {});
        written_ = 0;
        nextFrameEnd_ = static_cast<int64_t>(frame_) + maxDelay_;
        frames_ = 0;
        offset_ = 0;
    }
 
    void push(AudioBufferView<const T> reference, AudioBufferView<const T> delivery) noexcept
    {
        if (!prepared_) return;
        const int n = std::min(reference.getNumSamples(), delivery.getNumSamples());
        const int rc = reference.getNumChannels(), dc = delivery.getNumChannels();
        if (n <= 0 || rc <= 0 || dc <= 0) return;
        for (int i = 0; i < n; ++i)
        {
            double r = 0.0, d = 0.0;
            for (int c = 0; c < rc; ++c) r += static_cast<double>(reference.getChannel(c)[i]);
            for (int c = 0; c < dc; ++c) d += static_cast<double>(delivery.getChannel(c)[i]);
            const size_t slot = static_cast<size_t>(written_ % historyLength_);
            refHistory_[slot] = std::isfinite(r) ? r / rc : 0.0;
            delHistory_[slot] = std::isfinite(d) ? d / dc : 0.0;
            ++written_;
            if (written_ == nextFrameEnd_)
            {
                accumulateFrame();
                nextFrameEnd_ += hop_;
            }
        }
    }
 
    [[nodiscard]] Result estimate() noexcept
    {
        Result out;
        if (!prepared_ || frames_ < 1) return out;
        double lag = 0.0;
        bool inverted = false;
        if (!searchPeak(lag, inverted)) return out;
        const double sign = inverted ? -1.0 : 1.0;
 
        // Refine on the phase slope of the averaged cross-spectrum. The
        // delivery was framed offset_ late, so the spectrum holds the delay
        // minus offset_.
        double d = lag - static_cast<double>(offset_);
        for (int it = 0; it < 8; ++it)
        {
            double num = 0.0, den = 0.0;
            for (int k = 1; k < frame_ / 2; ++k)
            {
                const double w = 2.0 * pi<double> * k / frame_;
                const std::complex<double> s = cross_[static_cast<size_t>(k)] * sign
                    * std::polar(1.0, w * d);
                const double mag = std::abs(s);
                if (!(mag > 0.0)) continue;
                const double phi = std::arg(s);   // ~ -w * (delay - d)
                num += mag * w * phi;
                den += mag * w * w;
            }
            if (!(den > 0.0)) break;
            const double step = -num / den;
            d += step;
            if (std::abs(step) < 1e-9) break;
        }
 
        std::complex<double> agree {};
        double total = 0.0;
        for (int k = 1; k < frame_ / 2; ++k)
        {
            const double w = 2.0 * pi<double> * k / frame_;
            const std::complex<double> s = cross_[static_cast<size_t>(k)] * sign;
            agree += s * std::polar(1.0, w * d);
            total += std::abs(s);
        }
        out.delaySamples = d + static_cast<double>(offset_);
        out.delaySeconds = out.delaySamples / sampleRate_;
        out.confidence = (total > 0.0) ? std::clamp(std::abs(agree) / total, 0.0, 1.0) : 0.0;
        out.inverted = inverted;
        out.valid = total > 0.0;
        return out;
    }
 
    [[nodiscard]] int getFrameSize() const noexcept { return frame_; }
 
    static Result estimate(AudioBufferView<const T> reference, AudioBufferView<const T> delivery,
                           double sampleRate, int maxDelaySamples)
    {
        DelayEstimator est;
        est.prepare(sampleRate, maxDelaySamples);
        est.push(reference, delivery);
        // Flush: the last frames need the tail of the delivery.
        const int flush = est.frame_ + 2 * est.maxDelay_;
        std::vector<T> zeros(static_cast<size_t>(std::min(flush, 1 << 16)), T(0));
        const T* z = zeros.data();
        for (int left = flush; left > 0; left -= static_cast<int>(zeros.size()))
        {
            const int n = std::min(left, static_cast<int>(zeros.size()));
            est.push(AudioBufferView<const T>(&z, 1, n), AudioBufferView<const T>(&z, 1, n));
        }
        return est.estimate();
    }
 
private:
    void accumulateFrame() noexcept
    {
        // Frame the delivery at the whole-sample delay found so far (see the
        // file header); it settles within the first frames that carry
        // material and moves only if the evidence moves it.
        if (frames_ >= 1)
        {
            double lag = 0.0;
            bool inv = false;
            if (searchPeak(lag, inv))
            {
                const int newOffset = static_cast<int>(std::lround(lag));
                if (newOffset != offset_)
                {
                    // The spectrum so far was taken at the old offset: its
                    // phase holds (delay - old offset). Rotate it to the new.
                    for (int k = 0; k <= frame_ / 2; ++k)
                    {
                        const double w = 2.0 * pi<double> * k / frame_;
                        cross_[static_cast<size_t>(k)] *= std::polar(1.0, w * (newOffset - offset_));
                    }
                    offset_ = newOffset;
                }
            }
        }
 
        const int64_t refStart = written_ - maxDelay_ - frame_;
        const int64_t delStart = refStart + offset_;
        for (int i = 0; i < frame_; ++i)
        {
            const double w = window_[static_cast<size_t>(i)];
            bufA_[static_cast<size_t>(i)] = w * sample(refHistory_, refStart + i);
            bufB_[static_cast<size_t>(i)] = w * sample(delHistory_, delStart + i);
        }
        fft_->forward(bufA_.data(), bufA_.data());
        fft_->forward(bufB_.data(), bufB_.data());
        for (int k = 0; k <= frame_ / 2; ++k)
        {
            const std::complex<double> a(bufA_[static_cast<size_t>(2 * k)], bufA_[static_cast<size_t>(2 * k + 1)]);
            const std::complex<double> b(bufB_[static_cast<size_t>(2 * k)], bufB_[static_cast<size_t>(2 * k + 1)]);
            cross_[static_cast<size_t>(k)] += b * std::conj(a);
        }
        ++frames_;
    }
 
    [[nodiscard]] double sample(const std::vector<double>& h, int64_t index) const noexcept
    {
        if (index < 0 || index >= written_ || index < written_ - historyLength_) return 0.0;
        return h[static_cast<size_t>(index % historyLength_)];
    }
 
    bool searchPeak(double& lag, bool& inverted) noexcept
    {
        double peakMag = 0.0;
        for (const auto& s : cross_) peakMag = std::max(peakMag, std::abs(s));
        if (!(peakMag > 0.0)) return false;
        const double floor = peakMag * 1e-9;
        for (int k = 0; k <= frame_ / 2; ++k)
        {
            // Undo the framing offset so the correlation is in absolute lag.
            const double w = 2.0 * pi<double> * k / frame_;
            std::complex<double> s = cross_[static_cast<size_t>(k)] * std::polar(1.0, -w * offset_);
            const double m = std::abs(s);
            s = (m > floor) ? s / m : std::complex<double> {};
            corr_[static_cast<size_t>(2 * k)] = s.real();
            corr_[static_cast<size_t>(2 * k + 1)] = s.imag();
        }
        fft_->inverse(corr_.data(), corr_.data());
        auto at = [&](int tau) {
            const int m = ((tau % frame_) + frame_) % frame_;
            return corr_[static_cast<size_t>(m)];
        };
        int best = 0;
        double bestAbs = -1.0;
        for (int tau = -maxDelay_; tau <= maxDelay_; ++tau)
        {
            const double v = std::abs(at(tau));
            if (v > bestAbs) { bestAbs = v; best = tau; }
        }
        inverted = at(best) < 0.0;
        const double s = inverted ? -1.0 : 1.0;
        const double ym = s * at(best - 1), y0 = s * at(best), yp = s * at(best + 1);
        const double den = ym - 2.0 * y0 + yp;
        const double frac = (den < 0.0) ? std::clamp(0.5 * (ym - yp) / den, -0.5, 0.5) : 0.0;
        lag = static_cast<double>(best) + frac;
        return true;
    }
 
    double sampleRate_ = 48000.0;
    int maxDelay_ = 1;
    int frame_ = 4096;
    int hop_ = 2048;
    int historyLength_ = 0;
    bool prepared_ = false;
    std::unique_ptr<FFTReal<double>> fft_;
    std::vector<double> window_, refHistory_, delHistory_, bufA_, bufB_, corr_;
    std::vector<std::complex<double>> cross_;
    int64_t written_ = 0;
    int64_t nextFrameEnd_ = 0;
    int64_t frames_ = 0;
    int offset_ = 0;
};
 
} // namespace dspark