// SPDX-FileCopyrightText: 2026 0064ff
// SPDX-License-Identifier: MIT-0

#pragma once
#include <JuceHeader.h>

/** A processor to apply standard attack / release ballistics to an input signal. */
template <typename SampleType>
class MyBallisticsFilter
{
public:
    //==============================================================================
    MyBallisticsFilter()
    {
        // set defaults for variables
        setSampleRate(SampleType(44100));
        setSlew(SampleType(0.5));
        setTransientTime(SampleType(0.125));
        setAttackTime(SampleType(10));
        setReleaseTime(SampleType(100));
        setAdaptability(SampleType(2));
        setDetectionPower(SampleType(1));
        setLowCompOffset(SampleType(1));
        setHighCompOffset(SampleType(0));
        setCompressionAmountMultiplier(SampleType(0));
        setSag(SampleType(0));
    }

    //==============================================================================
    /** Sets the transient time in milliseconds. */
    void setTransientTime(SampleType newTransientTime)
    {
        if (transientTime == newTransientTime)
            return;

        transientTime = newTransientTime;
        updateTransientCache();
    }

    /** Sets the attack time in milliseconds. */
    void setAttackTime(SampleType newAttackTime)
    {
        if (attackTime == newAttackTime)
            return;

        attackTime = newAttackTime;
        updateAttackCache();
    }

    /** Sets the release time in milliseconds. */
    void setReleaseTime(SampleType newReleaseTime)
    {
        if (releaseTime == newReleaseTime)
            return;

        releaseTime = newReleaseTime;
        updateReleaseCache();
    }

    /** Sets slew as a percentage of the attack and release. Takes a value from 0 to 1, non-inclusive. */
    void setSlew(SampleType newSlew)
    {
        jassert(newSlew < SampleType(1));

        newSlew = std::min(newSlew, SampleType(.9999));

        if (slew == newSlew)
            return;

        slew = newSlew;
        updateSlewCache();
    }

    /** Sets the program-dependent multiplier for attack and release. */
    void setAdaptability(SampleType newAdaptability)
    {
        jassert(newAdaptability >= SampleType(1));

        newAdaptability = std::max(SampleType(1), newAdaptability);

        if (adaptability == newAdaptability)
            return;

        adaptability = newAdaptability;
        invAdaptability = SampleType(1) / newAdaptability;
    }

    /** Sets the power to raise the signal to. 1 = Peak. 2 = RMS. Expects a power that's 2^k, where k is an integer. */
    void setDetectionPower(SampleType newDetectionPower)
    {
        jassert(newDetectionPower > SampleType(0));
        jassert(std::abs(std::log2(newDetectionPower) - std::round(std::log2(newDetectionPower))) < 1e-6);

        detectionPower = newDetectionPower;
    }

    /** Sets the offset for the signal when minimal compression is occurring. 0 = symmetric (odd harmonics). 1 = asymmetric (even harmonics). */
    void setLowCompOffset(SampleType newLowCompOffset) {lowCompOffset=newLowCompOffset;}

    /** Sets the offset for the signal when maximal compression is occurring. 0 = symmetric (odd harmonics). 1 = asymmetric (even harmonics). */
    void setHighCompOffset(SampleType newHighCompOffset) {highCompOffset=newHighCompOffset;}

    /** Called by the compressor to let the filter know how much it's compressing. */
    void setCompressionAmountMultiplier(SampleType newCompressionAmountMultiplier) {compressionAmountMultiplier=newCompressionAmountMultiplier;}

    /** Called by the compressor to let the filter know how much sag to apply. */
    void setSag(SampleType newSag) {sag=newSag;}

    /** Sets values for testing. */
    void setTestValues(SampleType newTestValue1, SampleType newTestValue2, SampleType newTestValue3)
    {
        testValue1 = newTestValue1;
        testValue2 = newTestValue2;
        testValue3 = newTestValue3;
    }

    //==============================================================================
    /** Initialises the filter. */
    void prepare(juce::dsp::ProcessSpec const& spec)
    {
        jassert(spec.sampleRate > 0);
        jassert(spec.numChannels > 0);

        // resize vectors with channel count
        auto resize = [spec](auto&... vectors) {
            (vectors.resize(spec.numChannels), ...);
        };
        resize(envState, envSlowState, envPrevState, smoothedLevelState, smoothedTransientLevelState);

        // prepare with proper sample rate
        setSampleRate(spec.sampleRate);

        // caches need to be updated with proper sample rate
        updateTransientCache();
        updateAttackCache();
        updateReleaseCache();

        // initialize state vectors
        reset();
    }

    //==============================================================================
    /** Resets the internal state vectors of the filter. */
    void reset() {reset(0);}

    /** Resets the internal state vectors of the filter to the given initial value. */
    void reset(SampleType initialValue)
    {
        auto initialize = [&](auto&... vectors) {
            (std::fill(vectors.begin(), vectors.end(), initialValue), ...);
        };
        initialize(envState, envSlowState, envPrevState, smoothedLevelState, smoothedTransientLevelState);
    }

    //==============================================================================
    /** Processes a single sample. */
    SampleType processSample(int channel, SampleType inputValue) noexcept
    {
        jassert(isPositiveAndBelow(channel, envState.size()));

        // get values in vectors for processing
        SampleType env = envState[static_cast<std::size_t>(channel)];
        SampleType envSlow = envSlowState[static_cast<std::size_t>(channel)];
        SampleType envPrev = envPrevState[static_cast<std::size_t>(channel)];
        SampleType smoothedLevel = smoothedLevelState[static_cast<std::size_t>(channel)];
        SampleType smoothedTransientLevel = smoothedTransientLevelState[static_cast<std::size_t>(channel)];

        //-----------------------------------------------------------------------------------
        // level calculation

        // apply rectification and offset
        SampleType const sign = (inputValue > SampleType(0)) ? 1 : -1;
        SampleType const offsetDiff = highCompOffset - lowCompOffset;
        SampleType const effectiveOffset = lowCompOffset + offsetDiff * compressionAmountMultiplier;
        inputValue = std::abs(inputValue) * (SampleType(1) + effectiveOffset * sign);

        // apply detection power
        inputValue = quickPow(inputValue, detectionPower);

        //-----------------------------------------------------------------------------------
        // program-dependent attack and release

        // smooth level over attack time
        smoothedLevel = onePoleSmoother(smoothedLevel, inputValue, attackCoeff);
        smoothedLevel = std::max(smoothedLevel, SampleType(1e-6));

        // take ratio of instantaneous sample to smoothed level
        SampleType instantaneousRatio = inputValue / smoothedLevel;
        instantaneousRatio = std::max(instantaneousRatio, SampleType(1e-6));

        // attack and release get multipliers depending on the instantaneous ratio
        // these are limited by adaptability to mitigate extreme values
        SampleType instantaneousRatioInv = SampleType(1) / instantaneousRatio;
        instantaneousRatioInv = std::clamp(instantaneousRatioInv, invAdaptability, adaptability);
        SampleType attackMultiplier = instantaneousRatioInv;
        SampleType releaseMultiplier = instantaneousRatioInv;

        // smooth transient level
        // this is done after instantaneous ratio calculation so we can still catch instantaneous peaks
        smoothedTransientLevel = onePoleSmoother(smoothedTransientLevel, inputValue, transientCoeff);

        // attack reacts to overshoot of smoothed transient
        SampleType attackOvershootMultiplier = SampleType(1);
        if (smoothedTransientLevel > env)
        {
            SampleType const overshoot = smoothedTransientLevel - env;
            attackOvershootMultiplier = SampleType(1) / (SampleType(1)+overshoot/std::max(env, SampleType(1e-6)));
        }
        attackMultiplier *= attackOvershootMultiplier;

        // release reacts to undershoot of smoothed transient
        SampleType releaseUndershootMultiplier = SampleType(1);
        if (smoothedTransientLevel < env)
        {
            SampleType const undershoot = env - smoothedTransientLevel;
            releaseUndershootMultiplier = SampleType(1) / (SampleType(1)+undershoot/std::max(env, SampleType(1e-6)));
        }
        releaseMultiplier *= releaseUndershootMultiplier;

        // apply adaptability to final program-dependent multipliers
        attackMultiplier = std::clamp(attackMultiplier, invAdaptability, adaptability);
        releaseMultiplier = std::clamp(releaseMultiplier, invAdaptability, adaptability);

        //-----------------------------------------------------------------------------------
        // sag

        // sag isn't adaptability-limited because it's a byproduct of compression amount
        attackMultiplier *= SampleType(1) + sag;
        releaseMultiplier *= SampleType(1) + sag;

        //-----------------------------------------------------------------------------------
        // calculate attack and release coefficients

        // apply multipliers to envelope coefficients
        SampleType effectiveAttackCoeff = quickScaleCoeff(attackCoeff, attackMultiplier);
        SampleType effectiveReleaseCoeff = quickScaleCoeff(releaseCoeff, releaseMultiplier);

        // clamp to stable range
        effectiveAttackCoeff = std::clamp(effectiveAttackCoeff, SampleType(1e-6), SampleType(0.999999));
        effectiveReleaseCoeff = std::clamp(effectiveReleaseCoeff, SampleType(1e-6), SampleType(0.999999));

        //-----------------------------------------------------------------------------------
        // dual envelopes

        // calculate normal envelope
        SampleType const envCoeff = (smoothedTransientLevel > env) ? effectiveAttackCoeff : effectiveReleaseCoeff;
        env = onePoleSmoother(env, smoothedTransientLevel, envCoeff);

        // calculate slow envelope (double release time)
        SampleType const envCoeffSlow = (smoothedTransientLevel > envSlow)
            ? effectiveAttackCoeff
            : quickScaleCoeff(effectiveReleaseCoeff, SampleType(2));
        envSlow = onePoleSmoother(envSlow, smoothedTransientLevel, envCoeffSlow);

        // normal envelope follows slow envelope
        if (env < envSlow)
            env = onePoleSmoother(env, envSlow, effectiveReleaseCoeff);

        // calculate final result, with more compression resulting in slow-envelope dominance
        SampleType envFinal = env *  (1-compressionAmountMultiplier) +
                              envSlow * compressionAmountMultiplier;

        //-----------------------------------------------------------------------------------
        // slew rate limiting

        SampleType envDiff = envFinal - envPrev;

        // determine if slew should be for attack or release and use appropriate multiplier
        envDiff = (envDiff > SampleType(0)) ? std::min(envDiff, baseSlewAttack/attackMultiplier)
                                            : std::max(envDiff, -baseSlewRelease/releaseMultiplier);

        // apply slew-limited envelope
        envFinal = envPrev + envDiff;

        jassert(std::isfinite(envFinal));

        envPrev = envFinal;

        //-----------------------------------------------------------------------------------

        // update vectors
        envState[static_cast<std::size_t>(channel)] = env;
        envSlowState[static_cast<std::size_t>(channel)] = envSlow;
        envPrevState[static_cast<std::size_t>(channel)] = envPrev;
        smoothedLevelState[static_cast<std::size_t>(channel)] = smoothedLevel;
        smoothedTransientLevelState[static_cast<std::size_t>(channel)] = smoothedTransientLevel;

        // invert detection power
        return quickPowInv(envFinal, detectionPower);
    }

    //==============================================================================

    /** Quick power calculation. Expects a power that's 2^k, where k is an integer. Expects a non-negative value. */
    static SampleType quickPow(SampleType value, SampleType power) noexcept
    {
        jassert(value >= SampleType(0));

        SampleType newValue;

        if (power < SampleType(1)) // for powers of less than one, keep taking the square root
        {
            newValue = value;
            int invPower = static_cast<int>(SampleType(1)/power);
            while (invPower > 1)
            {
                newValue = std::sqrt(newValue);
                invPower >>= 1;
            }
        }
        else // otherwise, keep squaring the value until the power is reached
        {
            newValue = value;
            int p = static_cast<int>(power);
            for (int i = 1; i < p; ++i)
                newValue *= value;
        }

        return newValue;
    }

    /** Quick inverse power calculation. Expects a power that's 2^k, where k is an integer. Expects a non-negative value. */
    static SampleType quickPowInv(SampleType value, SampleType power) noexcept
    {
        jassert(value >= SampleType(0));

        SampleType newValue;

        if (power < SampleType(1)) // undo square roots by squaring
        {
            newValue = SampleType(1);
            int invPower = static_cast<int>(SampleType(1)/power);
            for (int p = 0; p < invPower; ++p)
                newValue *= value;
            return newValue;
        }
        else // undo squaring by taking square roots
        {
            newValue = value;
            int p = static_cast<int>(power);
            while (p > 1)
            {
                newValue = std::sqrt(newValue);
                p >>= 1;
            }
        }

        return newValue;
    }

    //==============================================================================

    /** Converts milliseconds to step amount per sample for processing. */
    SampleType msToSampleStep(SampleType timeMs) const noexcept
    {
        return SampleType(1000) / (timeMs*sampleRate);
    }

    /** Converts milliseconds to coefficient for processing. */
    SampleType msToCoeff(SampleType timeMs) const noexcept
    {
        return std::exp(-msToSampleStep(timeMs));
    }

    /** Smooths using a one-pole low-pass filter. */
    static SampleType onePoleSmoother(SampleType state, SampleType target, SampleType coeff) noexcept
    {
        return target + coeff*(state-target);
    }

    /** Scales a coefficient's time constant using the linearization of coefficient^(1/scale) around coefficient=1. */
    static SampleType quickScaleCoeff(SampleType coeff, SampleType scale) noexcept
    {
        return SampleType(1) + (coeff-SampleType(1))/scale;
    }

    //==============================================================================
private:
    /** Sets sample rate. */
    void setSampleRate(SampleType newSampleRate) {sampleRate = newSampleRate;}

    //==============================================================================

    /** Caches expensive operations related to transient time. Use this in prepare() to bypass the early return in setTransientTime(). */
    void updateTransientCache() noexcept {transientCoeff = msToCoeff(transientTime);}

    /** Caches expensive operations related to attack time. Use this in prepare() to bypass the early return in setAttackTime(). */
    void updateAttackCache() noexcept
    {
        attackCoeff = msToCoeff(attackTime);
        updateBaseSlewAttack();
    }

    /** Caches expensive operations related to release time. Use this in prepare() to bypass the early return in setReleaseTime(). */
    void updateReleaseCache() noexcept
    {
        releaseCoeff = msToCoeff(releaseTime);
        updateBaseSlewRelease();
    }

    /** Caches expensive operations related to slew. */
    void updateSlewCache() noexcept
    {
        updateBaseSlewAttack();
        updateBaseSlewRelease();
    }

    /** Caches base attack slew. */
    void updateBaseSlewAttack() noexcept
    {
        jassert(slew < SampleType(1));

        baseSlewAttack = msToSampleStep(attackTime/(SampleType(1)-slew));
    }

    /** Caches base release slew. */
    void updateBaseSlewRelease() noexcept
    {
        jassert(slew < SampleType(1));

        baseSlewRelease = msToSampleStep(releaseTime/(SampleType(1)-slew));
    }

    //==============================================================================

    std::vector<SampleType> envState, envSlowState, envPrevState, smoothedLevelState, smoothedTransientLevelState;

    static constexpr SampleType qNAN = std::numeric_limits<SampleType>::quiet_NaN();

    SampleType sampleRate=qNAN;

    SampleType transientCoeff=qNAN, attackCoeff=qNAN, releaseCoeff=qNAN;
    SampleType baseSlewAttack=qNAN, baseSlewRelease=qNAN;

    SampleType transientTime=qNAN, attackTime=qNAN, releaseTime=qNAN, adaptability=qNAN, invAdaptability=qNAN, slew=qNAN;
    SampleType detectionPower=qNAN, lowCompOffset=qNAN, highCompOffset=qNAN;
    SampleType compressionAmountMultiplier=qNAN, sag=qNAN;

    SampleType testValue1=qNAN, testValue2=qNAN, testValue3=qNAN;
};
