#!/usr/bin/env python3
"""MSM6258 (OKI/Dialogic) 4-bit ADPCM -- encoder, decoder, and the fact that
there are TWO decoders and they are not the same one.

The X68000's ADPCM is an OKI MSM6258V clocked at 8 MHz, dividing to 15,625 /
10,417 / 7,812.5 samples a second, 4 bits each, two samples to a byte
(FINDINGS 52, buscost.ADPCM_SAMPLE_HZ). The sample word is 12 bits signed.

WHY THIS FILE HAS TWO DECODERS. Nothing in this repo can be trusted to say what
the chip does, and the two references available on this machine DISAGREE:

  VARIANT 'shift'  delta = ((2*(n&7) + 1) * step) >> 3
                   This is ffmpeg's `adpcm_ima_oki`, and `gate_vs_ffmpeg()`
                   reproduces it SAMPLE-EXACT, so it is not a reading of source
                   code -- it is a measurement of the decoder that ships.

  VARIANT 'terms'  delta = step/8 + (n&4 ? step : 0) + (n&2 ? step/2 : 0)
                            + (n&1 ? step/4 : 0),  each term truncated
                   This is the OKI datasheet's own form, the one an ADPCM chip
                   can actually build out of shifts and adds, and it is what
                   MAME's okim6258 is understood to compute. NOT VERIFIED HERE:
                   no MAME source tree is on this machine (FINDINGS 64.4).

They differ on 445 of 2,268 sampled nibbles, by up to 4 in 12-bit units --
small, and small is not zero. Which one the machine runs is an open question
with an experiment attached: MAME's x68000 HAS an okim6258, so it can be asked
rather than argued about.

Nibble order is HIGH NIBBLE FIRST within a byte -- measured, not assumed, by the
same gate: reading low-first mismatches ffmpeg on 1,728 of 2,268 samples.
"""

# The 49-entry OKI step table. floor(16 * 1.1**k) for k in 0..48 -- built rather
# than pasted, so a transcription slip is not one of the things that can be
# wrong here.
STEP = [int(16 * 1.1**k) for k in range(49)]

# The nibble magnitude's effect on the step index. Four quiet nibbles walk it
# down one, four loud ones walk it up by more.
INDEX_ADJUST = (-1, -1, -1, -1, 2, 4, 6, 8)

SAMPLE_MIN, SAMPLE_MAX = -2048, 2047     # the 12-bit DAC word

VARIANTS = ("shift", "terms")


def delta(nibble, step, variant):
    """The reconstruction step for one nibble, in 12-bit units."""
    if variant == "shift":
        d = ((2 * (nibble & 7) + 1) * step) >> 3
    elif variant == "terms":
        d = step // 8
        if nibble & 4: d += step
        if nibble & 2: d += step // 2
        if nibble & 1: d += step // 4
    else:
        raise ValueError(f"unknown variant {variant!r}")
    return -d if nibble & 8 else d


def decode(nibbles, variant="shift"):
    """Nibbles -> 12-bit signed samples. State is (signal, step index), both
    zero at the start of a stream, which is what the chip resets to."""
    signal, idx, out = 0, 0, []
    for n in nibbles:
        signal += delta(n, STEP[idx], variant)
        signal = SAMPLE_MIN if signal < SAMPLE_MIN else (
                 SAMPLE_MAX if signal > SAMPLE_MAX else signal)
        idx += INDEX_ADJUST[n & 7]
        idx = 0 if idx < 0 else (48 if idx > 48 else idx)
        out.append(signal)
    return out


def encode(samples, variant="shift"):
    """12-bit signed samples -> nibbles.

    The nibble is chosen by EXHAUSTIVE SEARCH over all sixteen, minimising the
    reconstruction error of this sample. That is greedy rather than optimal --
    a nibble also moves the step index, so a locally worse choice can pay later
    -- but it is what a chip-matched encoder is expected to do and it costs
    nothing offline. The decoder is run INSIDE the loop, so the encoder can
    never drift away from what the decoder will reconstruct.
    """
    signal, idx, out = 0, 0, bytearray()
    for s in samples:
        step = STEP[idx]
        best, best_err = 0, None
        for n in range(16):
            v = signal + delta(n, step, variant)
            v = SAMPLE_MIN if v < SAMPLE_MIN else (SAMPLE_MAX if v > SAMPLE_MAX else v)
            err = (v - s) ** 2
            if best_err is None or err < best_err:
                best, best_err = n, err
        signal += delta(best, step, variant)
        signal = SAMPLE_MIN if signal < SAMPLE_MIN else (
                 SAMPLE_MAX if signal > SAMPLE_MAX else signal)
        idx += INDEX_ADJUST[best & 7]
        idx = 0 if idx < 0 else (48 if idx > 48 else idx)
        out.append(best)
    return bytes(out)


def pack(nibbles):
    """Nibbles -> bytes, HIGH NIBBLE FIRST. An odd count pads with a 0 nibble,
    which is the quietest one the format has (delta = step/8)."""
    n = bytes(nibbles)
    if len(n) & 1:
        n += b"\0"
    return bytes((n[i] << 4) | n[i + 1] for i in range(0, len(n), 2))


def unpack(data, count=None):
    out = bytearray()
    for b in data:
        out.append(b >> 4)
        out.append(b & 15)
    return bytes(out[:count] if count is not None else out)