SoundFlix/FractalVisualizer/Audio/AudioAnalyzer.cs

using System;
using System.Diagnostics;
using System.Numerics;
using System.Threading;
using System.Threading.Tasks;
using FractalVisualizer.Models;
using MathNet.Numerics.IntegralTransforms;

namespace FractalVisualizer.Audio;

public sealed class AudioAnalyzer : IDisposable
{
    private readonly AudioCapture _capture;
    private readonly int _fftSize;
    private readonly int _spectrumBins;
    private readonly float[] _hannWindow;
    private readonly float[] _timeSamples;
    private readonly Complex[] _fftBuffer;
    private readonly float[] _spectrumAccumulator;
    private readonly int[] _spectrumAccumulatorCounts;
    private readonly Stopwatch _clock = new();

    private readonly object _lifecycleSync = new();

    private CancellationTokenSource? _cts;
    private Task? _analysisTask;
    private int _sampleRate;

    private float _smoothedLevel;
    private float _smoothedBass;
    private float _smoothedMid;
    private float _smoothedHigh;
    private float _smoothedBeat;
    private float _energyBaseline;

    private Snapshot _latestSnapshot;
    private bool _disposed;

    public AudioAnalyzer(AudioCapture capture, int fftSize = 2048, int spectrumBins = 64)
    {
        if (fftSize <= 0 || (fftSize & (fftSize - 1)) != 0)
        {
            throw new ArgumentOutOfRangeException(nameof(fftSize), "FFT size must be a power of 2.");
        }

        if (spectrumBins <= 0)
        {
            throw new ArgumentOutOfRangeException(nameof(spectrumBins));
        }

        _capture = capture ?? throw new ArgumentNullException(nameof(capture));
        _fftSize = fftSize;
        _spectrumBins = spectrumBins;
        _sampleRate = capture.SampleRate;

        _hannWindow = BuildHannWindow(fftSize);
        _timeSamples = new float[fftSize];
        _fftBuffer = new Complex[fftSize];
        _spectrumAccumulator = new float[spectrumBins];
        _spectrumAccumulatorCounts = new int[spectrumBins];
        _latestSnapshot = new Snapshot(AudioFeatures.CreateEmpty(spectrumBins));

        _capture.SampleRateChanged += OnSampleRateChanged;
    }

    public int FftSize => _fftSize;

    public int SpectrumBins => _spectrumBins;

    public void Start()
    {
        ThrowIfDisposed();

        lock (_lifecycleSync)
        {
            if (_analysisTask is not null)
            {
                return;
            }

            _clock.Restart();
            _cts = new CancellationTokenSource();
            _analysisTask = Task.Run(() => AnalysisLoop(_cts.Token), _cts.Token);
        }
    }

    public void Stop()
    {
        Task? taskToWait;
        CancellationTokenSource? ctsToDispose;

        lock (_lifecycleSync)
        {
            if (_analysisTask is null || _cts is null)
            {
                return;
            }

            taskToWait = _analysisTask;
            ctsToDispose = _cts;
            _analysisTask = null;
            _cts = null;
        }

        ctsToDispose.Cancel();
        try
        {
            taskToWait.Wait();
        }
        catch (AggregateException ex) when (ex.InnerException is OperationCanceledException)
        {
        }
        finally
        {
            ctsToDispose.Dispose();
        }
    }

    public AudioFeatures GetLatestFeatures()
    {
        return Volatile.Read(ref _latestSnapshot).Features;
    }

    private void AnalysisLoop(CancellationToken cancellationToken)
    {
        while (!cancellationToken.IsCancellationRequested)
        {
            if (!_capture.RingBuffer.TryRead(_timeSamples))
            {
                Thread.Sleep(2);
                continue;
            }

            ProcessWindow();
        }
    }

    private void ProcessWindow()
    {
        int sampleRate = Volatile.Read(ref _sampleRate);
        int half = _fftSize / 2;

        for (int i = 0; i < _fftSize; i++)
        {
            _fftBuffer[i] = new Complex(_timeSamples[i] * _hannWindow[i], 0.0);
        }

        Fourier.Forward(_fftBuffer, FourierOptions.Matlab);

        Array.Clear(_spectrumAccumulator);
        Array.Clear(_spectrumAccumulatorCounts);

        float rmsAccumulator = 0f;
        for (int i = 0; i < _fftSize; i++)
        {
            float sample = _timeSamples[i];
            rmsAccumulator += sample * sample;
        }

        float rms = MathF.Sqrt(rmsAccumulator / _fftSize);

        float bassEnergy = 0f;
        float midEnergy = 0f;
        float highEnergy = 0f;

        for (int bin = 1; bin < half; bin++)
        {
            float frequency = (sampleRate * bin) / (float)_fftSize;
            float magnitude = (float)_fftBuffer[bin].Magnitude;

            int spectrumIndex = (int)((bin / (float)half) * _spectrumBins);
            if (spectrumIndex >= _spectrumBins)
            {
                spectrumIndex = _spectrumBins - 1;
            }

            _spectrumAccumulator[spectrumIndex] += magnitude;
            _spectrumAccumulatorCounts[spectrumIndex]++;

            if (frequency >= 20f && frequency < 200f)
            {
                bassEnergy += magnitude;
            }
            else if (frequency >= 200f && frequency < 2000f)
            {
                midEnergy += magnitude;
            }
            else if (frequency >= 2000f && frequency <= 16000f)
            {
                highEnergy += magnitude;
            }
        }

        float[] spectrum = new float[_spectrumBins];
        for (int i = 0; i < _spectrumBins; i++)
        {
            if (_spectrumAccumulatorCounts[i] == 0)
            {
                spectrum[i] = 0f;
                continue;
            }

            spectrum[i] = _spectrumAccumulator[i] / _spectrumAccumulatorCounts[i];
        }

        float instantEnergy = rms * rms;
        _energyBaseline = Lerp(_energyBaseline, instantEnergy, 0.02f);
        float beat = 0f;
        if (_energyBaseline > 1e-7f && instantEnergy > _energyBaseline * 1.5f)
        {
            beat = Math.Clamp((instantEnergy / _energyBaseline) - 1f, 0f, 1f);
        }

        _smoothedLevel = Smooth(_smoothedLevel, rms, 0.45f, 0.08f);
        _smoothedBass = Smooth(_smoothedBass, bassEnergy, 0.40f, 0.10f);
        _smoothedMid = Smooth(_smoothedMid, midEnergy, 0.30f, 0.08f);
        _smoothedHigh = Smooth(_smoothedHigh, highEnergy, 0.30f, 0.08f);
        _smoothedBeat = MathF.Max(beat, _smoothedBeat * 0.85f);

        AudioFeatures features = new()
        {
            Level = _smoothedLevel,
            Bass = _smoothedBass,
            Mid = _smoothedMid,
            High = _smoothedHigh,
            Beat = _smoothedBeat,
            Spectrum = spectrum,
            TimeSeconds = (float)_clock.Elapsed.TotalSeconds
        };

        Volatile.Write(ref _latestSnapshot, new Snapshot(features));
    }

    private void OnSampleRateChanged(int sampleRate)
    {
        Volatile.Write(ref _sampleRate, sampleRate);
    }

    private static float[] BuildHannWindow(int length)
    {
        float[] window = new float[length];
        if (length == 1)
        {
            window[0] = 1f;
            return window;
        }

        for (int i = 0; i < length; i++)
        {
            window[i] = 0.5f * (1f - MathF.Cos((2f * MathF.PI * i) / (length - 1)));
        }

        return window;
    }

    private static float Smooth(float current, float target, float attack, float release)
    {
        float factor = target > current ? attack : release;
        return current + ((target - current) * factor);
    }

    private static float Lerp(float a, float b, float t)
    {
        return a + ((b - a) * t);
    }

    private void ThrowIfDisposed()
    {
        if (_disposed)
        {
            throw new ObjectDisposedException(nameof(AudioAnalyzer));
        }
    }

    public void Dispose()
    {
        if (_disposed)
        {
            return;
        }

        Stop();
        _capture.SampleRateChanged -= OnSampleRateChanged;
        _disposed = true;
    }

    private sealed class Snapshot(AudioFeatures features)
    {
        public AudioFeatures Features { get; } = features;
    }
}