AutoGain.h

// DSPark - Professional Audio DSP Framework
// Copyright (c) 2026 Cristian Moresi - MIT License
 
#pragma once
 
#include "../Core/AudioBuffer.h"
#include "../Core/AudioSpec.h"
#include "../Core/Biquad.h"
#include "../Core/DspMath.h"
#include "../Core/StateBlob.h"
 
#include <algorithm>
#include <atomic>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <vector>
 
namespace dspark {
 
template <FloatType T>
class AutoGain
{
    // Ensure the atomic type won't trigger a hidden mutex lock in the audio thread
    static_assert(std::atomic<T>::is_always_lock_free,
        "AutoGain requires a lock-free float type for thread safety in the audio path.");
 
public:
    enum class Weighting
    {
        KWeighted, 
        Flat       
    };
 
    void prepare(const AudioSpec& spec)
    {
        if (!spec.isValid()) return; // release-safe: keep previous state
 
        sampleRate_ = spec.sampleRate;
        numChannels_ = spec.numChannels;
        shelf_ = BiquadCoeffs::makeKWeightingShelf(sampleRate_);
        highPass_ = BiquadCoeffs::makeKWeightingHighPass(sampleRate_);
        refFilters_.assign(static_cast<size_t>(numChannels_), WeightingState {});
        outFilters_.assign(static_cast<size_t>(numChannels_), WeightingState {});
        reset();
    }
 
    void pushReference(AudioBufferView<T> buffer) noexcept
    {
        if (std::min(buffer.getNumChannels(), numChannels_) <= 0 ||
            buffer.getNumSamples() <= 0)
            return;
        integrate(refMeanSquare_, buffer, refFilters_);
    }
 
    void compensate(AudioBufferView<T> buffer) noexcept
    {
        const int numCh = std::min(buffer.getNumChannels(), numChannels_);
        const int numSamples = buffer.getNumSamples();
 
        if (numSamples == 0 || numCh == 0) return;
 
        integrate(outMeanSquare_, buffer, outFilters_);
        const T refLevelDb = meanSquareToDb(refMeanSquare_);
        const T outLevelDb = meanSquareToDb(outMeanSquare_);
        T targetDb = refLevelDb - outLevelDb;
 
        // Clamp to safety limits
        const T maxComp = maxCompensation_.load(std::memory_order_relaxed);
        targetDb = std::clamp(targetDb, -maxComp, maxComp);
 
        // Silence bypass (-90 dB threshold)
        if (refLevelDb < SILENCE_THRESH_DB && outLevelDb < SILENCE_THRESH_DB)
            targetDb = T(0);
 
        // Calculate analytical end-state of the one-pole filter for the current block size:
        // alpha = exp(-N / (Fs * tau))
        const T smoothSecs = smoothTimeSecs_.load(std::memory_order_relaxed);
        const T alpha = static_cast<T>(std::exp(-static_cast<double>(numSamples)
                                  / (sampleRate_ * static_cast<double>(smoothSecs))));
        const T endCompensationDb = targetDb + (compensationDb_ - targetDb) * alpha;
 
        // Convert dB to linear gain for interpolation
        const T startGain = decibelsToGain(compensationDb_);
        const T endGain = decibelsToGain(endCompensationDb);
        const T gainStep = (endGain - startGain) / static_cast<T>(numSamples);
 
        // Apply linearly interpolated gain.
        // This loop structure guarantees no loop-carried dependencies, enabling strict SIMD vectorization.
        for (int ch = 0; ch < numCh; ++ch)
        {
            T* data = buffer.getChannel(ch);
            for (int i = 0; i < numSamples; ++i)
            {
                data[i] *= (startGain + static_cast<T>(i) * gainStep);
            }
        }
 
        // Update internal state for the next block
        compensationDb_ = endCompensationDb;
        // Publish it for cross-thread metering. The working member stays
        // audio-thread-private: getCompensationDb() used to read it directly,
        // which is a plain word written by the audio thread and read by
        // another -- a data race, not merely an approximate number.
        publishedCompensationDb_.store(compensationDb_, std::memory_order_relaxed);
    }
 
    void reset() noexcept
    {
        refMeanSquare_ = 0.0;
        outMeanSquare_ = 0.0;
        for (auto& f : refFilters_) f = {};
        for (auto& f : outFilters_) f = {};
        compensationDb_ = T(0);
        publishedCompensationDb_.store(T(0), std::memory_order_relaxed);
    }
 
    void setWeighting(Weighting w) noexcept
    {
        const int v = std::clamp(static_cast<int>(w), 0, static_cast<int>(Weighting::Flat));
        weighting_.store(static_cast<Weighting>(v), std::memory_order_relaxed);
    }
 
    [[nodiscard]] Weighting getWeighting() const noexcept
    {
        return weighting_.load(std::memory_order_relaxed);
    }
 
    [[nodiscard]] T getCompensationDb() const noexcept
    {
        return publishedCompensationDb_.load(std::memory_order_relaxed);
    }
 
    void setMaxCompensation(T dB) noexcept
    {
        if (!std::isfinite(dB)) return;
        maxCompensation_.store(std::abs(dB), std::memory_order_relaxed);
    }
 
    void setSmoothingTime(T ms) noexcept
    {
        if (!std::isfinite(ms)) return;
        smoothTimeSecs_.store(std::max<T>(ms * T(0.001), T(0.001)),
                              std::memory_order_relaxed);
    }
 
    [[nodiscard]] T getMaxCompensation() const noexcept
    {
        return maxCompensation_.load(std::memory_order_relaxed);
    }
 
    [[nodiscard]] T getSmoothingTime() const noexcept
    {
        return smoothTimeSecs_.load(std::memory_order_relaxed) * T(1000);
    }
 
    [[nodiscard]] std::vector<uint8_t> getState() const
    {
        // The blob stores float (setState reads float back); the explicit
        // casts also keep this overload resolvable when T is double.
        StateWriter w(stateId("AGAN"), 1);
        w.write("maxComp", static_cast<float>(maxCompensation_.load(std::memory_order_relaxed)));
        w.write("smoothMs", static_cast<float>(getSmoothingTime()));
        w.write("weighting", static_cast<int32_t>(getWeighting()));
        return w.blob();
    }
 
    bool setState(const uint8_t* data, size_t size)
    {
        StateReader r(data, size);
        if (!r.isValid() || r.processorId() != stateId("AGAN")) return false;
        setMaxCompensation(static_cast<T>(r.read("maxComp", 12.0f)));
        setSmoothingTime(static_cast<T>(r.read("smoothMs", 100.0f)));
        setWeighting(static_cast<Weighting>(r.read("weighting", 0)));
        return true;
    }
 
private:
    struct WeightingState { double s1a = 0.0, s2a = 0.0, s1b = 0.0, s2b = 0.0; };
 
    static inline double tdf2(double x, const BiquadCoeffs& c, double& s1, double& s2) noexcept
    {
        const double y = c.b0 * x + s1;
        s1 = c.b1 * x - c.a1 * y + s2;
        s2 = c.b2 * x - c.a2 * y;
        return y;
    }
 
    void integrate(double& meanSquare, AudioBufferView<T> buffer,
                   std::vector<WeightingState>& filters) noexcept
    {
        const int numCh = std::min({ buffer.getNumChannels(), numChannels_,
                                     static_cast<int>(filters.size()) });
        const int numSamples = buffer.getNumSamples();
        if (numCh <= 0 || numSamples <= 0) return;
 
        const bool weighted = weighting_.load(std::memory_order_relaxed) == Weighting::KWeighted;
        double sumSq = 0.0;
        for (int ch = 0; ch < numCh; ++ch)
        {
            const T* x = buffer.getChannel(ch);
            WeightingState& f = filters[static_cast<size_t>(ch)];
            if (weighted)
            {
                for (int i = 0; i < numSamples; ++i)
                {
                    const double y = tdf2(tdf2(static_cast<double>(x[i]), shelf_, f.s1a, f.s2a),
                                          highPass_, f.s1b, f.s2b);
                    sumSq += y * y;
                }
            }
            else
            {
                for (int i = 0; i < numSamples; ++i)
                    sumSq += static_cast<double>(x[i]) * static_cast<double>(x[i]);
            }
        }
 
        const double blockMs = sumSq / static_cast<double>(numSamples * numCh);
        if (!std::isfinite(blockMs))
        {
            for (auto& state : filters) state = {};
            return;
        }
        const double a = std::exp(-static_cast<double>(numSamples) / (sampleRate_ * kIntegrationSeconds));
        meanSquare = blockMs + (meanSquare - blockMs) * a;
    }
 
    [[nodiscard]] static T meanSquareToDb(double meanSquare) noexcept
    {
        return static_cast<T>(10.0 * std::log10(std::max(meanSquare, 1e-15)));
    }
 
    static constexpr double kIntegrationSeconds = 0.4; 
    static constexpr T SILENCE_THRESH_DB = T(-90);   
 
    double sampleRate_ = 44100.0;                    
    int numChannels_ = 0;                            
 
    std::atomic<T> smoothTimeSecs_{ T(0.100) };      
    std::atomic<Weighting> weighting_{ Weighting::KWeighted };
    double refMeanSquare_ = 0.0;                     
    double outMeanSquare_ = 0.0;                     
    BiquadCoeffs shelf_ = BiquadCoeffs::makeKWeightingShelf(48000.0);
    BiquadCoeffs highPass_ = BiquadCoeffs::makeKWeightingHighPass(48000.0);
    std::vector<WeightingState> refFilters_;
    std::vector<WeightingState> outFilters_;
    T compensationDb_ = T(0);                        
    std::atomic<T> publishedCompensationDb_{ T(0) }; 
 
    std::atomic<T> maxCompensation_{ T(12) };        
};
 
} // namespace dspark