"""What a streaming codec actually removes, measured by nulling against the source.

Encoders delay the signal (Vorbis and AAC prime, Opus has a pre-skip), so a naive
subtraction compares a file with a shifted copy of itself and reports noise. This
aligns first, by cross-correlation, then subtracts.

The control matters more than the results: the same pipeline is run on a lossless
round trip, and if that does not come back at around -300 dB the method is broken
and every other number here is decoration.

The codec settings come from projects.audio.CODEC_SPECS, so this measures exactly
what the product plays, not a separate guess about it.

    python scripts/codec_null_test.py master.wav
    python scripts/codec_null_test.py            # generated broadband source

The audio never leaves the machine it runs on, and nothing but numbers comes out:
residual in dB, per-band figures, the true-peak change. Run it where the file
already is. In this repo the dependencies live in the container, so:

    cp master.wav media/
    docker compose exec web uv run --directory /app python \
        /app/src/../scripts/codec_null_test.py /app/media/master.wav

A sine is useless here: a codec has almost nothing to throw away. Use real music,
or the generated source below, which is broadband with transients.
"""

import os
import shutil
import subprocess
import sys
import tempfile

import numpy as np
import soundfile as sf

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))


def specs():
    try:
        from projects.audio import CODEC_SPECS
        return CODEC_SPECS
    except Exception:
        # Standalone copy for anyone running this outside the repo. Keep in step
        # with projects/audio.py.
        return {
            'spotify': {'encoder': 'libvorbis', 'bitrate': '96k', 'ext': 'ogg',
                        'label': 'Ogg Vorbis 96 kbps'},
            'apple': {'encoder': 'aac', 'bitrate': '256k', 'ext': 'm4a',
                      'label': 'AAC-LC 256 kbps'},
            'youtube': {'encoder': 'libopus', 'bitrate': '128k', 'ext': 'opus',
                        'label': 'Opus 128 kbps'},
        }


def broadband_source(path, seconds=20, sr=44100):
    """Pink-ish noise with transients: something a codec has to make choices about."""
    rng = np.random.default_rng(7)
    n = seconds * sr
    white = rng.standard_normal(n)
    # One-pole cascade gives a rough pink tilt without scipy.
    pink = np.zeros(n)
    b = [0.0, 0.0, 0.0]
    for i in range(n):
        b[0] = 0.99765 * b[0] + white[i] * 0.0990460
        b[1] = 0.96300 * b[1] + white[i] * 0.2965164
        b[2] = 0.57000 * b[2] + white[i] * 1.0526913
        pink[i] = b[0] + b[1] + b[2] + white[i] * 0.1848
    pink /= np.max(np.abs(pink))
    # Transients every half second: what makes a codec's pre-echo audible.
    for t in range(0, seconds * 2):
        i = int(t * sr / 2)
        pink[i:i + 64] += np.hanning(64) * 0.9
    pink = np.clip(pink * 0.5, -1.0, 1.0)
    sf.write(path, np.column_stack([pink, pink]), sr, subtype='PCM_24')
    return path


def run(cmd):
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode != 0:
        raise RuntimeError((r.stderr or 'ffmpeg failed')[:400])


def to_wav(src, dst, sr):
    run(['ffmpeg', '-y', '-hide_banner', '-loglevel', 'error', '-i', src,
         '-ar', str(sr), '-ac', '2', '-c:a', 'pcm_f32le', dst])


def encode(src, dst, spec):
    cmd = ['ffmpeg', '-y', '-hide_banner', '-loglevel', 'error', '-i', src,
           '-map', '0:a:0', '-c:a', spec['encoder'], '-b:a', spec['bitrate']]
    if spec['encoder'] == 'aac':
        cmd += ['-profile:a', 'aac_low']
    run(cmd + [dst])


def align(ref, test, search=8192):
    """Integer-sample offset that best lines test up with ref.

    Encoder delay is the whole reason this is here: without it the residual is
    the difference between a signal and a shifted copy of itself, which is much
    larger than what the codec did.
    """
    n = min(len(ref), len(test), 400000)
    a = ref[:n] - ref[:n].mean()
    b = test[:n] - test[:n].mean()
    best, best_lag = -1e30, 0
    for lag in range(-search, search + 1, 1):
        if lag >= 0:
            x, y = a[:n - lag], b[lag:n]
        else:
            x, y = a[-lag:n], b[:n + lag]
        if len(x) < 1000:
            continue
        c = float(np.dot(x, y))
        if c > best:
            best, best_lag = c, lag
    return best_lag


def db(x):
    return -np.inf if x <= 0 else 20 * np.log10(x)


def rms(x):
    return float(np.sqrt(np.mean(x ** 2))) if len(x) else 0.0


def true_peak_ebur128(path):
    """True peak from ffmpeg's ebur128, which is a standard implementation and not
    ours. The numbers in the article come from here rather than from the estimate
    below: this is the claim most likely to be argued with, so it should not rest
    on a filter we wrote."""
    r = subprocess.run(
        ['ffmpeg', '-hide_banner', '-nostats', '-i', path,
         '-af', 'ebur128=peak=true', '-f', 'null', '-'],
        capture_output=True, text=True)
    lines = (r.stderr or '').splitlines()
    for i, line in enumerate(lines):
        if 'True peak' in line:
            for nxt in lines[i + 1:i + 4]:
                if 'Peak:' in nxt:
                    try:
                        return float(nxt.split('Peak:')[1].split('dBFS')[0])
                    except ValueError:
                        return None
    return None


def true_peak(x, sr, oversample=4):
    """Fallback: 4x linearly interpolated peak. It under-reads against ebur128 by
    a few tenths, which is the safe direction, and it exists only so the script
    still says something if ebur128 is unavailable."""
    n = len(x)
    up = np.interp(np.linspace(0, n - 1, n * oversample), np.arange(n), x)
    return db(float(np.max(np.abs(up))))


def band_residual(ref, diff, sr, edges=(0, 1000, 6000, 12000, 18000, 22050)):
    """Residual per band, relative to the source's energy in that band.

    The raw residual energy is unreadable: a band that carries most of the music
    will show a big number simply for being loud. What is wanted is how much of
    each band survived, so every figure is against that band's own source energy.
    """
    fr = np.fft.rfft(ref)
    fd = np.fft.rfft(diff)
    freqs = np.fft.rfftfreq(len(ref), 1 / sr)
    out = []
    for lo, hi in zip(edges, edges[1:]):
        m = (freqs >= lo) & (freqs < hi)
        if not m.any():
            continue
        er = float(np.sqrt(np.mean(np.abs(fr[m]) ** 2)))
        ed = float(np.sqrt(np.mean(np.abs(fd[m]) ** 2)))
        out.append(((lo, hi), db(ed / er) if er else -np.inf))
    return out


def measure(src_path, label, spec, tmp):
    ref, sr = sf.read(src_path, always_2d=True, dtype='float64')
    ref = ref[:, 0]

    enc = os.path.join(tmp, 'e.' + spec['ext'])
    dec = os.path.join(tmp, 'd.wav')
    encode(src_path, enc, spec)
    to_wav(enc, dec, sr)
    test, _ = sf.read(dec, always_2d=True, dtype='float64')
    test = test[:, 0]

    lag = align(ref, test)
    if lag >= 0:
        a, b = ref[:len(ref) - lag], test[lag:lag + len(ref) - lag]
    else:
        a, b = ref[-lag:], test[:len(ref) + lag]
    n = min(len(a), len(b))
    a, b = a[:n], b[:n]

    # Level: a codec can come back a hair off, and an uncorrected gain error
    # would show up as codec damage. Report both so nobody has to trust one.
    g = float(np.dot(a, b) / np.dot(b, b)) if np.dot(b, b) else 1.0
    diff_raw = a - b
    diff_gain = a - b * g

    return {
        'label': spec['label'],
        'lag': lag,
        'gain_db': db(abs(g)),
        'residual_db': db(rms(diff_raw) / rms(a)),
        'residual_gain_matched_db': db(rms(diff_gain) / rms(a)),
        'tp_src': true_peak_ebur128(src_path),
        'tp_enc': true_peak_ebur128(dec),
        'tp_src_est': true_peak(a, sr),
        'tp_enc_est': true_peak(b, sr),
        'bands': band_residual(a, diff_gain, sr),
        'bytes': os.path.getsize(enc),
    }


def check_alignment():
    """Shift a signal by a known amount and see whether align() finds it.

    Reporting "+0 samples" for every codec is exactly what a broken aligner looks
    like, so the aligner is asked to prove itself before any codec is measured.
    """
    rng = np.random.default_rng(1)
    x = rng.standard_normal(200000)
    for want in (0, 137, -412, 2048):
        y = np.roll(x, want)
        got = align(x, y, search=4096)
        if got != want:
            sys.exit('alignment is broken: shifted %+d, found %+d' % (want, got))
    return True


def main():
    if not shutil.which('ffmpeg'):
        sys.exit('ffmpeg not found')
    check_alignment()
    tmp = tempfile.mkdtemp()
    src = sys.argv[1] if len(sys.argv) > 1 else broadband_source(
        os.path.join(tmp, 'source.wav'))
    synthetic = len(sys.argv) <= 1
    print('source: %s (%.1f MB)%s\n' % (
        os.path.basename(src), os.path.getsize(src) / 1e6,
        '   [generated: verifies the method, do not publish these figures]'
        if synthetic else ''))

    # The control. If this is not far below everything else, stop reading.
    ctrl = measure(src, 'control', {'encoder': 'flac', 'bitrate': '0',
                                    'ext': 'flac', 'label': 'FLAC (lossless)'}, tmp)
    c = ctrl['residual_gain_matched_db']
    print('%-22s residual %s   <- the method returns "identical" when it should'
          % (ctrl['label'], 'perfect null' if c == -np.inf else '%.1f dB' % c))
    print()

    rows = []
    for key, spec in specs().items():
        r = measure(src, key, spec, tmp)
        rows.append((key, r))
        print('%-22s residual %8.1f dB  (aligned %+d samples, gain %+.2f dB)'
              % (r['label'], r['residual_gain_matched_db'], r['lag'], r['gain_db']))
        if r['tp_src'] is None or r['tp_enc'] is None:
            print('%-22s true peak %+.2f -> %+.2f dB (estimate)   size %.1f MB'
                  % ('', r['tp_src_est'], r['tp_enc_est'], r['bytes'] / 1e6))
        else:
            print('%-22s true peak %+.1f -> %+.1f dBTP (ebur128)   size %.1f MB'
                  % ('', r['tp_src'], r['tp_enc'], r['bytes'] / 1e6))
        parts = []
        for (lo, hi), rel in r['bands']:
            parts.append('%d-%dk %s' % (lo // 1000, hi // 1000,
                                        'gone' if rel > -0.5 else '%.0f dB' % rel))
        print('%-22s %s' % ('', '  '.join(parts)))
        print()

    if synthetic:
        print('The source above is noise with impulses, which is close to the worst\n'
              'case for a perceptual codec: noise is the first thing they discard.\n'
              'Run this on real music before quoting any of it.')
        print()
    if ctrl['residual_gain_matched_db'] > -200 and ctrl['residual_gain_matched_db'] != -np.inf:
        print('WARNING: the lossless control did not null. The numbers above mean '
              'nothing until it does.')


if __name__ == '__main__':
    main()