#!/usr/bin/env python3
"""Gate tools/encoder/adpcm.py against the only independent decoder on this
machine: ffmpeg's `adpcm_ima_oki`.

THIS FILE IS NOT ABOUT THE X68000's CHIP and after session 34 that distinction
is load-bearing. It checks this implementation against an independent one, so
its parameters stay ffmpeg's -- variant 'shift', high nibble first, a 12-bit
clamp, accumulator from 0. What the MACHINE's MSM6258 does is measured by
tools/bench/adpcm_run.sh and it is a different set of four values on all four
axes (adpcm.CHIP, FINDINGS 66). Do not "fix" the defaults here to match it: a
reference check whose reference has been adjusted to agree is not a check.

There is no ffmpeg ENCODER for this format -- `adpcm_ima_oki` is decode-only --
so the encoder here cannot be checked against a reference implementation. What
CAN be checked, and is, is that the decoder our encoder runs in its own loop is
byte-for-byte the decoder that ships in ffmpeg. An encoder that agrees with its
own wrong decoder is exactly the failure this catches.

Usage: verify_adpcm.py [wav_or_raw12 ...]
"""
import os, struct, subprocess, sys, random, math
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "encoder"))
import adpcm

TMP = "tmp/adpcm_gate"

# The step table as it is printed in the OKI datasheet and in every
# implementation of this format. adpcm.py BUILDS its table from 16*1.1**k; if
# the two ever disagree, one of them is a typo and this says which.
CANON = [16,17,19,21,23,25,28,31,34,37,41,45,50,55,60,66,73,80,88,97,107,118,
         130,143,157,173,190,209,230,253,279,307,337,371,408,449,494,544,598,
         658,724,796,876,963,1060,1166,1282,1411,1552]

fails = []
def ck(ok, msg):
    print(("OK  " if ok else "FAIL ") + msg)
    if not ok: fails.append(msg)


def ffmpeg_decode(data, rate=15625):
    """Decode packed OKI ADPCM through ffmpeg, by wrapping it in a WAV whose
    format tag is 0x0010 (WAVE_FORMAT_OKI_ADPCM). Returns 16-bit samples."""
    os.makedirs(TMP, exist_ok=True)
    fmt = struct.pack("<HHIIHHH", 0x0010, 1, rate, rate, 1, 4, 0)
    body = (b"WAVE" + b"fmt " + struct.pack("<I", len(fmt)) + fmt
                    + b"data" + struct.pack("<I", len(data)) + data)
    w = f"{TMP}/probe.wav"
    open(w, "wb").write(b"RIFF" + struct.pack("<I", len(body)) + body)
    raw = subprocess.check_output(
        ["ffmpeg", "-v", "error", "-i", w, "-f", "s16le", "-acodec", "pcm_s16le", "-"])
    return list(struct.unpack("<%dh" % (len(raw) // 2), raw))


def snr_db(ref, got):
    """Signal-to-noise over the 12-bit sample word."""
    num = sum(float(s) * s for s in ref)
    den = sum((float(a) - b) ** 2 for a, b in zip(ref, got))
    if den == 0: return float("inf")
    return 10.0 * math.log10(num / den) if num else float("-inf")


print("--- the step table ---")
ck(adpcm.STEP == CANON, f"49 entries, built = published (16*1.1**k), {adpcm.STEP[0]}..{adpcm.STEP[-1]}")

print("--- our decoder vs ffmpeg's adpcm_ima_oki ---")
random.seed(1234)
nibs = ([7]*12 + [i % 16 for i in range(256)]
        + [random.randrange(16) for _ in range(4000)])
data = adpcm.pack(nibs)
ff = ffmpeg_decode(data)
ours = [v * 16 for v in adpcm.decode(adpcm.unpack(data, len(nibs)), "shift")]
ck(len(ff) == len(ours), f"sample count {len(ff)} = {len(ours)}")
ck(ff == ours, f"variant 'shift' is SAMPLE-EXACT vs ffmpeg over {len(nibs)} nibbles")

# NEGATIVE CONTROL. A gate that passes whatever it is handed proves nothing;
# reading the nibbles the other way round has to go red, or "high nibble first"
# is an assertion rather than a measurement.
lowfirst = [adpcm.unpack(data, len(nibs))[i ^ 1] for i in range(len(nibs))]
bad = [v * 16 for v in adpcm.decode(lowfirst, "shift")]
ndiff = sum(1 for a, b in zip(ff, bad) if a != b)
ck(ndiff > 0, f"low-nibble-first DISAGREES on {ndiff}/{len(ff)} -- so what ffmpeg reads is measured, not assumed")
# AND IT IS A FACT ABOUT A FILE FORMAT, NOT ABOUT A CHIP. Session 33 recorded
# this line as "nibble order: HIGH FIRST, measured", which it is -- of the VOX
# convention ffmpeg implements. Session 34 asked the machine's own MSM6258 the
# same question through HD63450 channel 3 and got the OTHER answer: the chip
# takes the LOW nibble of a delivered byte first (FINDINGS 66), and encoding
# for the wrong one of the two costs -25.7 dB on the Singe window. The two
# claims do not conflict; they are about different things, and only one of them
# is about the machine this is being ported to.

print("--- and the second variant is not the same decoder ---")
terms = [v * 16 for v in adpcm.decode(adpcm.unpack(data, len(nibs)), "terms")]
d = [abs(a - b) // 16 for a, b in zip(ff, terms)]
nd = sum(1 for x in d if x)
ck(nd > 0, f"variant 'terms' differs on {nd}/{len(d)} samples, max {max(d)} in 12-bit units"
           " -- and 'terms' is the one the machine runs (adpcm_run.sh, FINDINGS 66)")

print("--- and getting the variant wrong is NOT a rounding error ---")
# THE MEASUREMENT THAT CHANGED THIS FROM A FOOTNOTE INTO AN OPEN ITEM. The two
# variants differ by at most 3 in 12-bit units PER SAMPLE, which reads like
# something nobody could hear. ADPCM is RECURSIVE -- the delta is added to a
# running predictor and the nibble also moves the step index -- so the
# disagreement does not stay where it happens. It is the same shape as the
# codec's temporal recursion (64.1), one dimension down.
if os.path.exists("tmp/au_singe.raw"):
    raw = open("tmp/au_singe.raw", "rb").read()
    pcm = struct.unpack("<%dh" % (len(raw) // 2), raw)
    ref = [max(-2048, min(2047, x >> 4)) for x in pcm]
    nib = adpcm.encode(ref, "shift")
    same, cross = adpcm.decode(nib, "shift"), adpcm.decode(nib, "terms")
    err = [abs(x - y) for x, y in zip(same, cross)]
    print(f"    encoded 'shift', decoded 'shift': SNR {snr_db(ref, same):6.2f} dB")
    print(f"    encoded 'shift', decoded 'terms': SNR {snr_db(ref, cross):6.2f} dB"
          f"   <- the noise is LOUDER THAN THE SIGNAL")
    print(f"    per-sample disagreement over {len(err):,} samples: max {max(err)}, "
          f"mean {sum(err)/len(err):.1f} in 12-bit units")
    ck(snr_db(ref, cross) < 0,
       "a 3-LSB formula disagreement costs ~25 dB, because ADPCM is RECURSIVE")
else:
    print("    SKIPPED: no tmp/au_singe.raw (tools/encoder/extract_audio.py)")

print("--- the encoder, through the gated decoder ---")
for path in (sys.argv[1:] or []):
    raw = open(path, "rb").read()
    if raw[:4] == b"RIFF":
        raw = subprocess.check_output(["ffmpeg", "-v", "error", "-i", path,
              "-f", "s16le", "-ac", "1", "-ar", "15625", "-"])
    pcm16 = struct.unpack("<%dh" % (len(raw) // 2), raw)
    src = [max(-2048, min(2047, s >> 4)) for s in pcm16]
    nib = adpcm.encode(src, "shift")
    packed = adpcm.pack(nib)
    rec_ff = [v // 16 for v in ffmpeg_decode(packed)][:len(src)]
    rec_us = adpcm.decode(nib, "shift")
    ck(rec_ff == rec_us,
       f"{os.path.basename(path)}: encoder's own reconstruction = ffmpeg's, {len(src)} samples")
    print(f"    {len(src)} samples, {len(packed)} B, SNR {snr_db(src, rec_us):.2f} dB "
          f"(12-bit word; the source is already quantised to it)")

print("ADPCM GATE " + ("GREEN" if not fails else f"RED: {len(fails)} failed"))
sys.exit(1 if fails else 0)