jRay/Jellyfin.Plugin.JRay/Services/AudioSignature.cs

using System;
using System.Collections.Generic;

namespace Jellyfin.Plugin.JRay.Services;

/// <summary>
/// The JRay v1 content-derived audio signature: spectral peak bins taken from
/// the centre of a media file, so a truth file is self-identifying.
/// </summary>
/// <remarks>
/// The construction is owned by the public server specification §3 and is
/// implemented a second time, in C++, by the extraction pipeline
/// (<c>src/audio_signature.*</c>, extraction <c>IR-004</c>). <b>The two must
/// agree byte for byte</b> — a signature that differs in any parameter simply
/// does not match, which defeats the entire point of having one. Every constant
/// below is therefore load-bearing, and any change to one is a change to the
/// <c>v1:</c> prefix as well.
/// <para>
/// The server specification's prose is not sufficient to reproduce a byte
/// stream, so the details it leaves open are pinned by the golden fixture
/// shared with the extraction repo (JR-043,
/// <c>fixtures/audio/jray_audio_v1_golden.json</c>), and restated here:
/// </para>
/// <list type="bullet">
/// <item>Arithmetic is IEEE-754 <c>double</c> throughout. <c>float</c> is not
/// sufficient: the fixture has frames whose two strongest bands are within 1.3%
/// of each other.</item>
/// <item>Samples arrive as FFmpeg's native <c>s16 -&gt; flt</c> conversion,
/// <c>x * (1/32768)</c>, widened to double here.</item>
/// <item>Whole frames only:
/// <c>n_frames = 1 + (n_samples - 4096) / 1024</c>, integer division.</item>
/// <item>Hann window, <b>periodic</b>: <c>0.5 * (1 - cos(2*pi*n/4096))</c>, not
/// the symmetric <c>N-1</c> variant.</item>
/// <item>Plain radix-2 FFT, no normalisation; magnitude is
/// <c>sqrt(re^2 + im^2)</c>.</item>
/// <item>A band's value is the <b>mean</b> of the linear magnitudes in it, so a
/// wide high band is not favoured over a narrow low one.</item>
/// <item>The peak is the <c>argmax</c> over the 32 bands, ties to the lowest
/// index. The specification's log is a monotone squash and so cannot change an
/// argmax; it is applied only where it is observable, in the energy class.</item>
/// </list>
/// <para>
/// The FFT is written out here rather than taken from a library for the same
/// reason the pipeline writes its own: it is a fixed, fully specified transform,
/// and a dependency whose version could change the numerics is a liability when
/// the output has to be identical across two languages.
/// </para>
/// </remarks>
// TRACES: JR-042 | SR-003
public static class AudioSignature
{
    /// <summary>Sample rate the signature is computed at, in Hz.</summary>
    public const int SampleRate = 11025;

    /// <summary>STFT frame size, in samples.</summary>
    public const int FrameSize = 4096;

    /// <summary>STFT hop size, in samples (~93 ms).</summary>
    public const int HopSize = 1024;

    /// <summary>Number of logarithmically spaced bands.</summary>
    public const int NumBands = 32;

    /// <summary>Low edge of the analysed band, in Hz.</summary>
    public const double BandLoHz = 300.0;

    /// <summary>High edge of the analysed band, in Hz.</summary>
    public const double BandHiHz = 3000.0;

    /// <summary>Length of the analysed window, in seconds.</summary>
    public const double WindowSec = 120.0;

    /// <summary>
    /// Length of the analysed window, in samples (120.000 s at 11025 Hz).
    /// </summary>
    /// <remarks>
    /// The decoded window is truncated to exactly this, so the frame count is
    /// the same for every input and does not wobble with seek granularity or a
    /// resampler tail.
    /// </remarks>
    public const int WindowSamples = 1323000;

    /// <summary>Frames a full window yields: <c>1 + (1323000 - 4096) / 1024</c>.</summary>
    /// <remarks>
    /// The server specification says "~1290" and accepts a tolerance; the exact
    /// count follows from the framing rule and is 1288.
    /// </remarks>
    public const int ExpectedFrames = 1288;

    /// <summary>Guard added to both sides of the energy ratio.</summary>
    public const double EnergyEps = 1e-12;

    /// <summary>
    /// The signature's own version prefix, separate from <c>schema_version</c>.
    /// </summary>
    /// <remarks>
    /// A future change to the DSP chain must be <i>detectable</i> rather than
    /// silently producing signatures that no longer match (JR-045).
    /// </remarks>
    public const string VersionPrefix = "v1:";

    // Class thresholds on log10(E_frame / E_median). They deliberately straddle
    // r = 0 rather than sit on it, so the median frame itself is never on a
    // boundary.
    private static readonly double[] EnergyClassEdges = [-0.6, -0.2, 0.2];

    private static readonly (int Low, int High)[] Bands = BuildBandTable();
    private static readonly double[] Window = BuildHannWindow();
    private static readonly int[] BitReversal = BuildBitReversal();
    private static readonly double[][] TwiddleReal = BuildTwiddles(cosine: true);
    private static readonly double[][] TwiddleImag = BuildTwiddles(cosine: false);

    /// <summary>
    /// Gets the half-open FFT bin range <c>[Low, High)</c> owned by each of the
    /// 32 log-spaced bands.
    /// </summary>
    /// <remarks>
    /// Exposed so the golden fixture can assert the table itself rather than
    /// only the signature it produces: the band table is the part of the
    /// construction most likely to drift between two implementations.
    /// </remarks>
    /// <returns>One range per band, contiguous and non-overlapping.</returns>
    public static IReadOnlyList<(int Low, int High)> BandFftBins() => Bands;

    /// <summary>
    /// Packs one byte per whole STFT frame: a 5-bit peak band index and a 2-bit
    /// energy class.
    /// </summary>
    /// <remarks>
    /// The byte layout is <c>(band &lt;&lt; 2) | class</c>, so bit 7 is always
    /// clear and an arbitrary byte is not a valid signature. That structural
    /// constraint is what the server validates on upload, and what keeps the
    /// field from being usable as a payload channel.
    /// </remarks>
    /// <param name="mono">Mono samples at <see cref="SampleRate"/>, in [-1, 1).</param>
    /// <returns>One byte per frame; empty when not even one frame fits.</returns>
    public static byte[] PackFrames(ReadOnlySpan<float> mono)
    {
        if (mono.Length < FrameSize)
        {
            return [];
        }

        var frames = 1 + ((mono.Length - FrameSize) / HopSize);
        var binLow = Bands[0].Low;
        var binHigh = Bands[NumBands - 1].High;         // exclusive
        double binCount = binHigh - binLow;

        var re = new double[FrameSize];
        var im = new double[FrameSize];
        var peak = new byte[frames];
        var energy = new double[frames];

        for (var f = 0; f < frames; f++)
        {
            var src = mono.Slice(f * HopSize, FrameSize);
            for (var n = 0; n < FrameSize; n++)
            {
                re[n] = src[n] * Window[n];
                im[n] = 0.0;
            }

            Fft(re, im);

            // Per-band mean magnitude. The bands tile 300-3000 Hz with no gaps
            // and no overlaps, so the frame's band-limited energy is the sum of
            // the band sums — accumulated in band order, because the order of a
            // floating-point summation is part of the contract.
            var best = -1.0;
            var bestBand = 0;
            var total = 0.0;
            for (var b = 0; b < NumBands; b++)
            {
                var (low, high) = Bands[b];
                var sum = 0.0;
                for (var k = low; k < high; k++)
                {
                    sum += Math.Sqrt((re[k] * re[k]) + (im[k] * im[k]));
                }

                total += sum;
                var mean = sum / (high - low);
                if (mean > best)
                {
                    best = mean;                        // ties -> lowest index
                    bestBand = b;
                }
            }

            peak[f] = (byte)bestBand;
            energy[f] = total / binCount;
        }

        // The reference is the upper median of the frame energies: an actually
        // observed value rather than the average of the two middle ones, so it
        // is bit-reproducible. It is also gain-invariant — loudness
        // normalisation must not change a signature — and barely moves when the
        // window is trimmed.
        var sorted = (double[])energy.Clone();
        Array.Sort(sorted);
        var reference = sorted[sorted.Length / 2];

        var packed = new byte[frames];
        for (var f = 0; f < frames; f++)
        {
            var r = Math.Log10((energy[f] + EnergyEps) / (reference + EnergyEps));
            packed[f] = (byte)(((peak[f] & 0x1F) << 2) | (EnergyClass(r) & 0x03));
        }

        return packed;
    }

    /// <summary>
    /// Computes the full signature string for a decoded centre window.
    /// </summary>
    /// <param name="mono">Mono samples at <see cref="SampleRate"/>, in [-1, 1).</param>
    /// <returns>
    /// <c>v1:</c> followed by the base64 of <see cref="PackFrames"/>, or
    /// <c>null</c> when not even one frame fits.
    /// </returns>
    public static string? FromMonoSamples(ReadOnlySpan<float> mono)
    {
        var packed = PackFrames(mono);
        if (packed.Length == 0)
        {
            return null;
        }

        // Standard alphabet with '=' padding, which is what the server and the
        // pipeline both emit.
        return VersionPrefix + Convert.ToBase64String(packed);
    }

    private static int EnergyClass(double ratio)
    {
        if (ratio < EnergyClassEdges[0])
        {
            return 0;
        }

        if (ratio < EnergyClassEdges[1])
        {
            return 1;
        }

        if (ratio < EnergyClassEdges[2])
        {
            return 2;
        }

        return 3;
    }

    // edge[b] = 300 * (3000/300)^(b/32); band b owns FFT bins
    // [k_lo[b], k_lo[b+1]) with k_lo[b] = ceil(edge[b] / hz_per_bin). Taking the
    // ceiling once, into an integer table, means membership is never decided by
    // a float comparison per bin per frame — which is where two implementations
    // would otherwise be free to disagree.
    private static (int Low, int High)[] BuildBandTable()
    {
        var hzPerBin = (double)SampleRate / FrameSize;
        var edges = new int[NumBands + 1];
        for (var b = 0; b <= NumBands; b++)
        {
            var hz = BandLoHz * Math.Pow(BandHiHz / BandLoHz, (double)b / NumBands);
            edges[b] = (int)Math.Ceiling(hz / hzPerBin);
        }

        var table = new (int Low, int High)[NumBands];
        for (var b = 0; b < NumBands; b++)
        {
            table[b] = (edges[b], edges[b + 1]);
        }

        return table;
    }

    private static double[] BuildHannWindow()
    {
        var w = new double[FrameSize];
        for (var n = 0; n < FrameSize; n++)
        {
            w[n] = 0.5 * (1.0 - Math.Cos(2.0 * Math.PI * n / FrameSize));
        }

        return w;
    }

    private static int[] BuildBitReversal()
    {
        var bits = 0;
        while ((1 << bits) < FrameSize)
        {
            bits++;
        }

        var rev = new int[FrameSize];
        for (var i = 0; i < FrameSize; i++)
        {
            var r = 0;
            for (var b = 0; b < bits; b++)
            {
                if ((i & (1 << b)) != 0)
                {
                    r |= 1 << (bits - 1 - b);
                }
            }

            rev[i] = r;
        }

        return rev;
    }

    // Twiddles are precomputed per stage from cos/sin of -2*pi*j/len, so the
    // angle is an exactly reproducible double in either language and only the
    // library's own rounding of cos/sin (<= 1 ulp) can differ — orders of
    // magnitude below the decision margins the golden fixture records.
    private static double[][] BuildTwiddles(bool cosine)
    {
        var stages = new List<double[]>();
        for (var len = 2; len <= FrameSize; len <<= 1)
        {
            var half = len / 2;
            var stage = new double[half];
            for (var j = 0; j < half; j++)
            {
                var angle = -2.0 * Math.PI * j / len;
                stage[j] = cosine ? Math.Cos(angle) : Math.Sin(angle);
            }

            stages.Add(stage);
        }

        return [.. stages];
    }

    // Radix-2 decimation-in-time complex FFT, in place, no normalisation.
    private static void Fft(double[] re, double[] im)
    {
        for (var i = 0; i < FrameSize; i++)
        {
            var j = BitReversal[i];
            if (i < j)
            {
                (re[i], re[j]) = (re[j], re[i]);
                (im[i], im[j]) = (im[j], im[i]);
            }
        }

        var stage = 0;
        for (var len = 2; len <= FrameSize; len <<= 1, stage++)
        {
            var half = len / 2;
            var wr = TwiddleReal[stage];
            var wi = TwiddleImag[stage];
            for (var start = 0; start < FrameSize; start += len)
            {
                for (var j = 0; j < half; j++)
                {
                    var a = start + j;
                    var b = a + half;
                    var tr = (re[b] * wr[j]) - (im[b] * wi[j]);
                    var ti = (re[b] * wi[j]) + (im[b] * wr[j]);
                    re[b] = re[a] - tr;
                    im[b] = im[a] - ti;
                    re[a] += tr;
                    im[a] += ti;
                }
            }
        }
    }
}