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 float[][] _spectrumRefsByMode;
    private readonly float[] _levelRefsByMode = new float[2];
    private readonly float[] _bassRefsByMode = new float[2];
    private readonly float[] _midRefsByMode = new float[2];
    private readonly float[] _highRefsByMode = new float[2];
    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 float _lastChordChangeTimeSeconds;
    private int _stableNoteIndex = -1;
    private float _stableNoteStrength;
    private float _lastNoteChangeTimeSeconds;

    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];
        _spectrumRefsByMode =
        [
            new float[spectrumBins],
            new float[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()
    {
        AudioInputMode mode = _capture.CurrentMode;
        SensitivityProfile profile = mode == AudioInputMode.Microphone ? SensitivityProfile.Microphone : SensitivityProfile.SystemOutput;
        int modeIndex = mode == AudioInputMode.Microphone ? 0 : 1;
        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 >= 120f && frequency <= 6000f && 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;
            }

            float raw = _spectrumAccumulator[i] / _spectrumAccumulatorCounts[i];
            spectrum[i] = NormalizeWithReference(raw, ref _spectrumRefsByMode[modeIndex][i], profile.SpectrumGain, profile.NoiseFloor, profile.ReferenceRise, profile.ReferenceDecay);
        }

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

        float levelNorm = NormalizeWithReference(rms, ref _levelRefsByMode[modeIndex], profile.LevelGain, profile.NoiseFloor, profile.ReferenceRise, profile.ReferenceDecay);
        float bassNorm = NormalizeWithReference(bassEnergy, ref _bassRefsByMode[modeIndex], profile.BassGain, profile.NoiseFloor, profile.ReferenceRise, profile.ReferenceDecay);
        float midNorm = NormalizeWithReference(midEnergy, ref _midRefsByMode[modeIndex], profile.MidGain, profile.NoiseFloor, profile.ReferenceRise, profile.ReferenceDecay);
        float highNorm = NormalizeWithReference(highEnergy, ref _highRefsByMode[modeIndex], profile.HighGain, profile.NoiseFloor, profile.ReferenceRise, profile.ReferenceDecay);

        _smoothedLevel = Smooth(_smoothedLevel, levelNorm, 0.45f, 0.08f);
        _smoothedBass = Smooth(_smoothedBass, bassNorm, 0.40f, 0.10f);
        _smoothedMid = Smooth(_smoothedMid, midNorm, 0.30f, 0.08f);
        _smoothedHigh = Smooth(_smoothedHigh, highNorm, 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 whistleBoost = GetWhistleBoost(dominantFrequency);
        float dominantStrength = Clamp01(dominantMagnitude * profile.NoteStrengthGain * (1f + whistleBoost * 0.65f));
        DetectChord(_chroma, out int chordRoot, out int chordType, out float chordConfidence);
        ApplyDetectionStability(
            now,
            profile,
            noteIndex,
            dominantStrength,
            chordRoot,
            chordType,
            chordConfidence,
            out int stableNoteIndex,
            out float stableNoteStrength,
            out int stableChordRoot,
            out int stableChordType,
            out float stableChordConfidence);

        AudioFeatures features = new()
        {
            Level = _smoothedLevel,
            Bass = _smoothedBass,
            Mid = _smoothedMid,
            High = _smoothedHigh,
            Beat = _smoothedBeat,
            Bpm = _smoothedBpm,
            DominantFrequencyHz = dominantFrequency,
            DetectedNoteIndex = stableNoteIndex,
            DetectedNoteStrength = stableNoteStrength,
            DetectedChordRootIndex = stableChordRoot,
            DetectedChordType = stableChordType,
            DetectedChordConfidence = stableChordConfidence,
            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 static float GetWhistleBoost(float frequencyHz)
    {
        if (frequencyHz <= 0f)
        {
            return 0f;
        }

        // Human whistles usually sit around ~1kHz-4kHz, with strong note clarity in this band.
        if (frequencyHz < 800f)
        {
            return 0f;
        }

        if (frequencyHz > 5000f)
        {
            return 0f;
        }

        float centered = (frequencyHz - 800f) / (5000f - 800f);
        return Clamp01(centered);
    }

    private static float NormalizeWithReference(float value, ref float reference, float gain, float noiseFloor, float referenceRise, float referenceDecay)
    {
        float v = MathF.Max(0f, value - noiseFloor);
        if (v > reference)
        {
            reference = Lerp(reference, v, referenceRise);
        }
        else
        {
            reference = Lerp(reference, v, referenceDecay);
        }

        if (reference < 1e-6f)
        {
            reference = 1e-6f;
        }

        return Clamp01((v / reference) * gain);
    }

    private void ApplyDetectionStability(
        float now,
        SensitivityProfile profile,
        int noteIndex,
        float noteStrength,
        int chordRoot,
        int chordType,
        float chordConfidence,
        out int stableNoteIndex,
        out float stableNoteStrength,
        out int stableChordRoot,
        out int stableChordType,
        out float stableChordConfidence)
    {
        stableNoteIndex = _stableNoteIndex;
        stableNoteStrength = _stableNoteStrength * 0.96f;

        if (noteIndex >= 0 && noteStrength >= profile.NoteConfidenceThreshold)
        {
            bool changed = _stableNoteIndex != noteIndex;
            bool canSwitch = (now - _lastNoteChangeTimeSeconds) >= profile.SwitchHoldSeconds;
            if (!changed || canSwitch || noteStrength > _stableNoteStrength * 1.25f)
            {
                stableNoteIndex = noteIndex;
                stableNoteStrength = noteStrength;
                if (changed)
                {
                    _lastNoteChangeTimeSeconds = now;
                }
            }
        }

        _stableNoteIndex = stableNoteIndex;
        _stableNoteStrength = stableNoteStrength;

        stableChordRoot = _lastChordRoot;
        stableChordType = _lastChordType;
        stableChordConfidence = 0f;
        if (stableChordRoot >= 0)
        {
            stableChordConfidence = Clamp01(chordConfidence * 0.85f);
        }

        if (chordRoot >= 0 && chordConfidence >= profile.ChordConfidenceThreshold)
        {
            bool changed = _lastChordRoot != chordRoot || _lastChordType != chordType;
            bool canSwitch = (now - _lastChordChangeTimeSeconds) >= profile.SwitchHoldSeconds;
            if (!changed || canSwitch || chordConfidence > stableChordConfidence * 1.2f)
            {
                _lastChordRoot = chordRoot;
                _lastChordType = chordType;
                if (changed)
                {
                    _lastChordChangeTimeSeconds = now;
                }
            }

            stableChordRoot = _lastChordRoot;
            stableChordType = _lastChordType;
            stableChordConfidence = chordConfidence;
        }
    }

    private readonly struct SensitivityProfile(
        float levelGain,
        float bassGain,
        float midGain,
        float highGain,
        float spectrumGain,
        float noteStrengthGain,
        float noiseFloor,
        float referenceRise,
        float referenceDecay,
        float switchHoldSeconds,
        float noteConfidenceThreshold,
        float chordConfidenceThreshold)
    {
        public float LevelGain { get; } = levelGain;
        public float BassGain { get; } = bassGain;
        public float MidGain { get; } = midGain;
        public float HighGain { get; } = highGain;
        public float SpectrumGain { get; } = spectrumGain;
        public float NoteStrengthGain { get; } = noteStrengthGain;
        public float NoiseFloor { get; } = noiseFloor;
        public float ReferenceRise { get; } = referenceRise;
        public float ReferenceDecay { get; } = referenceDecay;
        public float SwitchHoldSeconds { get; } = switchHoldSeconds;
        public float NoteConfidenceThreshold { get; } = noteConfidenceThreshold;
        public float ChordConfidenceThreshold { get; } = chordConfidenceThreshold;

        public static SensitivityProfile Microphone { get; } = new(
            levelGain: 1.20f,
            bassGain: 1.25f,
            midGain: 1.20f,
            highGain: 1.15f,
            spectrumGain: 1.40f,
            noteStrengthGain: 0.036f,
            noiseFloor: 0.0009f,
            referenceRise: 0.16f,
            referenceDecay: 0.010f,
            switchHoldSeconds: 0.18f,
            noteConfidenceThreshold: 0.11f,
            chordConfidenceThreshold: 0.18f);

        public static SensitivityProfile SystemOutput { get; } = new(
            levelGain: 1.00f,
            bassGain: 1.00f,
            midGain: 1.00f,
            highGain: 1.00f,
            spectrumGain: 1.18f,
            noteStrengthGain: 0.022f,
            noiseFloor: 0.0004f,
            referenceRise: 0.20f,
            referenceDecay: 0.012f,
            switchHoldSeconds: 0.18f,
            noteConfidenceThreshold: 0.12f,
            chordConfidenceThreshold: 0.14f);
    }

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