SoundFlix/FractalVisualizer/Audio/AudioAnalyzer.cs

using System;
using System.Collections.Generic;
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 float[] _chroma = new float[12];
    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 float _smoothedBpm;
    private float _lastBeatTimeSeconds = -10f;
    private bool _beatGateOpen = true;
    private readonly Queue<float> _beatIntervals = new();
    private const int MaxBeatIntervals = 8;
    private int _lastChordRoot = -1;
    private int _lastChordType = -1;

    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;
        float dominantMagnitude = 0f;
        float dominantFrequency = 0f;
        Array.Clear(_chroma);

        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;
            }

            if (frequency >= 80f && frequency <= 2000f && magnitude > dominantMagnitude)
            {
                dominantMagnitude = magnitude;
                dominantFrequency = frequency;
            }

            if (frequency >= 60f && frequency <= 5000f)
            {
                int pitchClass = GetNoteIndexFromFrequency(frequency);
                if (pitchClass >= 0)
                {
                    _chroma[pitchClass] += 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);

        float now = (float)_clock.Elapsed.TotalSeconds;
        bool beatHit = beat > 0.62f && _beatGateOpen;
        if (beatHit)
        {
            float delta = now - _lastBeatTimeSeconds;
            if (delta > 0.22f && delta < 2.0f)
            {
                _beatIntervals.Enqueue(delta);
                while (_beatIntervals.Count > MaxBeatIntervals)
                {
                    _beatIntervals.Dequeue();
                }
            }

            _lastBeatTimeSeconds = now;
            _beatGateOpen = false;
        }
        else if (beat < 0.20f)
        {
            _beatGateOpen = true;
        }

        float bpmInstant = CalculateBpm();
        _smoothedBpm = Smooth(_smoothedBpm, bpmInstant, 0.18f, 0.04f);

        int noteIndex = GetNoteIndexFromFrequency(dominantFrequency);
        float dominantStrength = Clamp01(dominantMagnitude * 0.02f);
        DetectChord(_chroma, out int chordRoot, out int chordType, out float chordConfidence);

        AudioFeatures features = new()
        {
            Level = _smoothedLevel,
            Bass = _smoothedBass,
            Mid = _smoothedMid,
            High = _smoothedHigh,
            Beat = _smoothedBeat,
            Bpm = _smoothedBpm,
            DominantFrequencyHz = dominantFrequency,
            DetectedNoteIndex = noteIndex,
            DetectedNoteStrength = dominantStrength,
            DetectedChordRootIndex = chordRoot,
            DetectedChordType = chordType,
            DetectedChordConfidence = chordConfidence,
            Spectrum = spectrum,
            TimeSeconds = now
        };

        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 float CalculateBpm()
    {
        if (_beatIntervals.Count == 0)
        {
            return 0f;
        }

        float sum = 0f;
        foreach (float interval in _beatIntervals)
        {
            sum += interval;
        }

        float avg = sum / _beatIntervals.Count;
        if (avg <= 1e-5f)
        {
            return 0f;
        }

        float bpm = 60f / avg;
        return Math.Clamp(bpm, 40f, 220f);
    }

    private static int GetNoteIndexFromFrequency(float frequencyHz)
    {
        if (frequencyHz <= 0f)
        {
            return -1;
        }

        float midi = 69f + (12f * MathF.Log2(frequencyHz / 440f));
        int note = ((int)MathF.Round(midi)) % 12;
        if (note < 0)
        {
            note += 12;
        }

        return note;
    }

    private void DetectChord(float[] chroma, out int root, out int type, out float confidence)
    {
        float total = 0f;
        for (int i = 0; i < chroma.Length; i++)
        {
            total += chroma[i];
        }

        if (total <= 1e-6f)
        {
            root = -1;
            type = -1;
            confidence = 0f;
            return;
        }

        float bestScore = 0f;
        int bestRoot = -1;
        int bestType = -1;
        for (int r = 0; r < 12; r++)
        {
            float major = ScoreChord(chroma, r, 0);
            if (major > bestScore)
            {
                bestScore = major;
                bestRoot = r;
                bestType = 0;
            }

            float minor = ScoreChord(chroma, r, 1);
            if (minor > bestScore)
            {
                bestScore = minor;
                bestRoot = r;
                bestType = 1;
            }
        }

        float conf = Clamp01(bestScore / (total + 1e-6f));
        if (conf < 0.17f)
        {
            root = -1;
            type = -1;
            confidence = conf;
            return;
        }

        if (_lastChordRoot >= 0 && _lastChordType >= 0)
        {
            float held = ScoreChord(chroma, _lastChordRoot, _lastChordType);
            if ((bestRoot != _lastChordRoot || bestType != _lastChordType) && held * 1.08f > bestScore)
            {
                bestRoot = _lastChordRoot;
                bestType = _lastChordType;
                bestScore = held;
                conf = Clamp01(bestScore / (total + 1e-6f));
            }
        }

        _lastChordRoot = bestRoot;
        _lastChordType = bestType;
        root = bestRoot;
        type = bestType;
        confidence = conf;
    }

    private static float ScoreChord(float[] chroma, int root, int type)
    {
        int third = type == 0 ? (root + 4) % 12 : (root + 3) % 12;
        int fifth = (root + 7) % 12;
        int second = (root + 2) % 12;
        int fourth = (root + 5) % 12;
        float score = chroma[root] * 1.0f + chroma[third] * 0.9f + chroma[fifth] * 0.8f;
        score -= (chroma[second] + chroma[fourth]) * 0.15f;
        return Math.Max(0f, score);
    }

    private static float Clamp01(float value)
    {
        if (value < 0f)
        {
            return 0f;
        }

        if (value > 1f)
        {
            return 1f;
        }

        return value;
    }

    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;
    }
}