ST_FFTThread.h

// Copyright (c) 2025 Aaron Trotter (ShaderTech). All Rights Reserved.
 
#pragma once
#include "CoreMinimal.h"
#include "HAL/ThreadSafeBool.h"
#include "HAL/RunnableThread.h"
#include "HAL/Runnable.h"
#include "HAL/Event.h"
 
#if WITH_KISSFFT
#include "kiss_fft.h"
#include "tools/kiss_fftr.h"
#endif
#include "Main/ST_AudioProcessingSettings.h"

class ST_AUDIOPROCESSING_API FST_FFTThread : public FRunnable
{
public:
    FST_FFTThread(
        FString InName,
        FFFTIOData InFFTIOData,
        FFFTConfig InFFTConfig,
        class UST_AudioIOBase* InAnalyserBase);

    ~FST_FFTThread();

    virtual bool CreateThread(bool bStartPaused);

    virtual void PauseThread();

    virtual void ContinueThread();

    virtual void Main();

    virtual bool IsThreadPaused() const;

    virtual void EnsureCompletion();
 
protected:
    virtual bool Init() override;

    virtual uint32 Run() override;

    virtual void Stop() override;

    virtual void Exit() override;
 
private:
    bool Compute_FFT();

    void Compute_FrequencyBands(uint32 ChannelIndex);

    void Compute_Amplitude(uint32 ChannelIndex);

    void Compute_Pitch(uint32 ChannelIndex);

    bool Compute_Beat(uint32 ChannelIndex);

    float PerformChecksum(const float* buffer, uint32 size);

    void Filter(FFFTFilterConfig InFilterConfig, TArray<float>& FilteredMagnitudes, uint32 ChannelIndex);

    void Filter_Butterworth(FFFTFilterConfig InFilterConfig, TArray<float>& FilteredMagnitudes, uint32 ChannelIndex);

    void Filter_ChebyshevTypeI(FFFTFilterConfig InFilterConfig, TArray<float>& FilteredMagnitudes, uint32 ChannelIndex);

    bool Windowing_Hann(uint32 ChannelIndex);

    bool Windowing_Hamming(uint32 ChannelIndex);

    bool Windowing_BlackmanHarris(uint32 ChannelIndex);

    bool Windowing_Gaussian(uint32 ChannelIndex);

    void SpectrumSpacing_Log();

    void SpectrumSpacing_OctaveSubdivision();

    void SpectrumSpacing_Linear();

    void SpectrumSpacing_Quadratic();
 
private:
    TUniquePtr<FRunnableThread> Thread = nullptr;

    FEvent* m_semaphore = nullptr;

    FThreadSafeBool m_Kill;

    FThreadSafeBool m_Pause;

    FString m_name = "";

    FCriticalSection FFTLock;

    FFFTIOData FFTIOData;

    FFFTConfig FFTConfig;

    TFunction<void(const TArray<float>&)> m_OnProcessedAudioReady;
 
#if WITH_KISSFFT
    kiss_fft_cpx* fftOutputBuffer = nullptr;

    kiss_fftr_cfg cfg = NULL;
#endif

    TArray<float> RawMagnitudes;

    TArray<float> RawMagnitudesHistory;

    TArray<float> FrequencyBandsHistory;

    TArray<float> FrequencyBands;

    TArray<float> FrequencyBandEdges;

    TArray<uint32> FrequencyBinsByBand;

    TArray<float> AmplitudeHistory;

    TArray<float> PitchHistory;

    TArray<float> BeatHistory;

    TArray<float> BeatRawMagnitudes;

    TArray<bool> BeatDetectedForChannel;

    TArray<bool> ProcessedChannel;

    class UST_AudioIOBase* cachedAnalyserBase = nullptr;

    TArray<float*> PaddedFFTInputBuffer;

    TArray<float> LastBufferChecksum;

    uint32 ChannelMultiplier = 1;
};
 
inline bool FST_FFTThread::IsThreadPaused() const
{
    return (bool)m_Pause;
}
 
inline void FST_FFTThread::PauseThread()
{
    m_Pause = true;
}