"""Momentary bass fatness of rendered audio, after Schmedeke's dissertation.

Model (L. Schmedeke, "BASS DROP! Bass in modern popular Dance Music",
Hamburg 2021, section 5.4.3, adjusted R² .773 on dance-music excerpts):

    l = 10 / (1 + 4.5487 · exp(0.13 · (62.9366 − SPL<200)))
    h = 10 / (1 + 1.4317 · exp(2.5575 · (0.3299 − HR<200)))
    Fat_mom = 0.14689 · l · h          (0 .. 14.7)

SPL<200 is the C-weighted sound-pressure level below 200 Hz in dB(C), HR<200
the MPEG-7 harmonic ratio of the same band, both over 50 ms windows, with the
material normalised to −23 LUFS first.

Assumption: the thesis does not give the dB(C) its reference tone played at.
`--offset` maps dBFS to dB(C); the default 100 dB puts −20 dBFS RMS at 80 dB(C),
a usual studio calibration. The thesis reports most complex excerpts at 70–80
dB(C) below 200 Hz; the summary prints where the material lands so the offset
can be checked.

Requires ffmpeg, numpy and matplotlib. Describes one model's view of the low
band; it is not a quality score.
"""
import os
os.environ.setdefault('MPLCONFIGDIR', '/tmp/fatness-matplotlib')
import argparse
import json
import subprocess
from pathlib import Path

import numpy as np

RATE = 48000
WINDOW_S = 0.05
HR_RATE = 4000
HR_MAX_LAG_S = 0.04


def load(path):
    raw = subprocess.check_output(['ffmpeg', '-v', 'error', '-i', str(path), '-f', 'f32le',
                                   '-ac', '2', '-ar', str(RATE), 'pipe:1'])
    return np.frombuffer(raw, dtype='<f4').reshape(-1, 2).T.astype(np.float64)


def k_weight_db(f):
    """Magnitude of the BS.1770 K-weighting (shelf + high-pass) at `f` Hz, in dB."""
    def biquad(b, a):
        z = np.exp(-2j * np.pi * f / RATE)
        return (b[0] + b[1] * z + b[2] * z * z) / (a[0] + a[1] * z + a[2] * z * z)
    shelf = biquad([1.53512485958697, -2.69169618940638, 1.19839281085285],
                   [1.0, -1.69065929318241, 0.73248077421585])
    high = biquad([1.0, -2.0, 1.0], [1.0, -1.99004745483398, 0.99007225036621])
    return 20 * np.log10(np.abs(shelf * high) + 1e-30)


def c_weight_db(f):
    """IEC 61672 C-weighting at `f` Hz, in dB."""
    f2 = np.maximum(f, 1e-3) ** 2
    r = (12194.0 ** 2 * f2) / ((f2 + 20.6 ** 2) * (f2 + 12194.0 ** 2))
    return 20 * np.log10(r) + 0.06


def filtered(x, gain_db):
    """Zero-phase filter by an FFT magnitude response; energy is what is measured."""
    n = x.shape[-1]
    size = 1 << (n - 1).bit_length()  # zero padding: fast FFT sizes only
    spec = np.fft.rfft(x, n=size, axis=-1)
    f = np.fft.rfftfreq(size, 1 / RATE)
    return np.fft.irfft(spec * 10 ** (gain_db(f) / 20), n=size, axis=-1)[..., :n]


def integrated_lufs(x):
    """BS.1770-4 integrated loudness with absolute and relative gates."""
    y = filtered(x, k_weight_db)
    block, step = int(0.4 * RATE), int(0.1 * RATE)
    sq = (y * y).sum(axis=0)
    c = np.concatenate([[0.], np.cumsum(sq)])
    starts = np.arange(0, len(sq) - block + 1, step)
    power = (c[starts + block] - c[starts]) / block
    loud = -0.691 + 10 * np.log10(power + 1e-30)
    keep = power[loud > -70]
    if keep.size == 0:
        return -np.inf
    relative = -0.691 + 10 * np.log10(keep.mean()) - 10
    keep = power[(loud > -70) & (loud > relative)]
    return -0.691 + 10 * np.log10(keep.mean())


def band_below(x, hz):
    return filtered(x, lambda f: np.where(f < hz, 0.0, -300.0))


def harmonic_ratio(mono, rate):
    """MPEG-7 harmonic ratio per 50 ms window: the largest normalised
    autocorrelation after the zero-lag lobe, with lags up to 40 ms."""
    n = int(WINDOW_S * rate)
    lags = int(HR_MAX_LAG_S * rate)
    out = []
    for start in range(lags, len(mono) - n + 1, n):
        s = mono[start:start + n]
        e0 = s @ s
        if e0 <= 1e-12:
            out.append(0.0)
            continue
        # Row m − 1 holds the window delayed by m samples.
        past = np.lib.stride_tricks.sliding_window_view(mono[start - lags:start - 1 + n], n)[::-1]
        gamma = (past @ s) / np.sqrt(e0 * np.einsum('ij,ij->i', past, past) + 1e-30)
        below = np.nonzero(gamma <= 0)[0]
        if below.size == 0:
            out.append(0.0)
            continue
        out.append(float(max(gamma[below[0]:].max(), 0.0)))
    return np.array(out), lags


def fatness(spl, hr):
    l = 10 / (1 + 4.5487 * np.exp(0.13 * (62.9366 - spl)))
    h = 10 / (1 + 1.4317 * np.exp(2.5575 * (0.3299 - hr)))
    return 0.14689 * l * h


def measure(path, offset):
    x = load(path)
    lufs = integrated_lufs(x)
    x = x * 10 ** ((-23 - lufs) / 20)
    low = band_below(x, 200.)
    mono = low.mean(axis=0)
    # HR on a 4 kHz copy; the band is below 200 Hz, so nothing aliases.
    hr, skip = harmonic_ratio(mono[::RATE // HR_RATE], HR_RATE)
    weighted = filtered(low, c_weight_db)
    n = int(WINDOW_S * RATE)
    first = skip * (RATE // HR_RATE)
    sq = (weighted * weighted).sum(axis=0) / 2
    frames = len(hr)
    spl = np.array([10 * np.log10(sq[first + i * n:first + (i + 1) * n].mean() + 1e-30)
                    for i in range(frames)]) + offset
    times = (first + np.arange(frames) * n) / RATE
    return lufs, times, spl, hr, fatness(spl, hr)


def main():
    parser = argparse.ArgumentParser(description=__doc__.split('\n\n')[0])
    parser.add_argument('files', nargs='*', type=Path)
    parser.add_argument('--folder', type=Path, default=Path('references/devath'))
    parser.add_argument('--offset', type=float, default=100.0,
                        help='dB(C) for a 0 dBFS RMS signal (default 100)')
    parser.add_argument('--output', type=Path, required=True,
                        help='folder for fatness.png and fatness.json')
    args = parser.parse_args()
    files = args.files or sorted(args.folder.glob('*.flac'))
    if not files:
        parser.error('no audio files')
    args.output.mkdir(parents=True, exist_ok=True)
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(len(files), 1, squeeze=False, figsize=(15, 2.2 * len(files)),
                             constrained_layout=True)
    summary = []
    for ax, path in zip(axes[:, 0], files):
        lufs, t, spl, hr, fat = measure(path, args.offset)
        smooth = np.convolve(fat, np.ones(8) / 8, mode='same')  # 400 ms
        pct = lambda v: {str(p): round(float(np.percentile(v, p)), 3) for p in (10, 50, 90)}
        row = {'file': path.name, 'lufs': round(float(lufs), 2), 'spl_c_below_200': pct(spl),
               'hr_below_200': pct(hr), 'fatness': pct(fat),
               'share_spl_70_80': round(float(np.mean((spl >= 70) & (spl <= 80))), 3)}
        summary.append(row)
        print(json.dumps(row))
        ax.plot(t, fat, color='#ccc', lw=0.5)
        ax.plot(t, smooth, color='#e0115f', lw=1.2)
        ax.set_ylim(0, 14.7)
        ax.set_xlim(0, t[-1] if len(t) else 1)
        ax.set_title(f'{path.stem}  ({lufs:.1f} LUFS, normalised to -23)', loc='left', fontsize=9)
        ax.set_ylabel('Fat_mom')
    axes[-1, 0].set_xlabel('seconds')
    fig.savefig(args.output / 'fatness.png', dpi=110)
    (args.output / 'fatness.json').write_text(json.dumps(
        {'offset_db': args.offset, 'window_s': WINDOW_S, 'tracks': summary}, indent=2))


if __name__ == '__main__':
    main()