#!/usr/bin/env python3
"""What a BRANCH costs the chip, and what an encoder could do about it.

    python3 tools/analysis/37_audio_seek.py [container.dlxp] [--raw au.raw]
    python3 tools/analysis/37_audio_seek.py --gate          # the check.sh stage

FINDINGS 71.  Session 39 put an audio seek path in src/player/packed.s and ran
it: 132,162 B across a real branch, every byte accounted for in MAME's own
capture.  That settled the BYTES.  This settles what is left, which is a
property of the codec rather than of the player and which no counter in the
player can reach.

THE MSM6258'S ACCUMULATOR HAS NO LEAKAGE TERM.  It is a pure integrator of
deltas, clamped, and nothing pulls it back toward zero.  So a branch that hands
the chip bytes chosen for a state it is not in does not produce a transient with
a time constant -- it produces a DC OFFSET THAT NEVER DECAYS.  The machine run
measures both designs at one branch point; this measures the CENSUS, over every
frame boundary of the container, and prices the only fix that is worth anything,
which is in the encoder and not in the player.

  * PLAY THROUGH the branch: the chip keeps whatever accumulator and step index
    the previous scene's audio left it in.  Unbounded, and its decay is the
    signal's own clamping rather than the recursion forgetting.

  * STOP and re-PLAY: the accumulator goes to the container's `init` and the
    step index to 0 -- a state this script knows exactly, so the error is
    EXACTLY `init - acc(target)`, constant, forever.

  * ...and the third option is the ENCODER'S: encode the stream with the
    predictor RESET at every point a branch can land on.  Then a re-PLAYing
    player is not close, it is exact.  What that costs is a codec question and
    is measured below.

NAME THE LAYER.  Everything here is host arithmetic over one container and its
source PCM.  The chip's four axes are the ones FINDINGS 66 measured on the
machine and 67.3 put in the header; the branch behaviour is the one session 39
ran.  No emulator is involved and no rate is claimed.
"""
import argparse, math, os, statistics, sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "encoder"))
import adpcm
from dlxp import DLXP


def acc_trajectory(nibbles, dec):
    """The decoder's accumulator after every sample.  This IS the encoder's
    assumed state, because adpcm.encode runs its decoder inside its own search
    loop -- the encoder cannot hold a state the decoder will not reach."""
    lo, hi = adpcm.clamp_bounds(dec["bits"])
    sig, idx = dec["init"], 0
    out = []
    for n in nibbles:
        sig += adpcm.delta(n, adpcm.STEP[idx], dec["variant"])
        sig = lo if sig < lo else (hi if sig > hi else sig)
        idx += adpcm.INDEX_ADJUST[n & 7]
        idx = 0 if idx < 0 else (48 if idx > 48 else idx)
        out.append((sig, idx))
    return out


def encode_reset(src12, dec, period):
    """adpcm.encode with the predictor forced back to (init, 0) every `period`
    samples.  period=None is the ordinary encode.

    THIS CHANGES THE BYTES, so it is a container property and not a flag a
    player can set.  It is written here rather than in tools/encoder/adpcm.py
    because nothing has decided to ship it: 71.5 is the trade and the deciding
    number is a hardware one."""
    if period is None:
        return adpcm.encode(src12, variant=dec["variant"], init=dec["init"],
                            bits=dec["bits"])
    out = bytearray()
    for i in range(0, len(src12), period):
        out += adpcm.encode(src12[i:i + period], variant=dec["variant"],
                            init=dec["init"], bits=dec["bits"])
    return bytes(out)


def snr(ref, got):
    n = min(len(ref), len(got))
    sig = sum(x * x for x in ref[:n])
    err = sum((ref[i] - got[i]) ** 2 for i in range(n))
    if err == 0:
        return float("inf")
    return 10 * math.log10(sig / err) if sig else float("-inf")


def decode_reset(nib, dec, period):
    if period is None:
        return list(adpcm.decode_state(nib, variant=dec["variant"],
                                       init=dec["init"], bits=dec["bits"])[0])
    out = []
    for i in range(0, len(nib), period):
        out += list(adpcm.decode_state(nib[i:i + period], variant=dec["variant"],
                                       init=dec["init"], bits=dec["bits"])[0])
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("container", nargs="?", default="tmp/packed_singe.dlxp")
    ap.add_argument("--raw", default="tmp/au_singe.raw")
    ap.add_argument("--gate", action="store_true")
    a = ap.parse_args()

    d = DLXP(a.container)
    if not d.has_audio:
        sys.exit("this container is silent -- there is no branch to price")
    dec = d.decoder()
    lo, hi = adpcm.clamp_bounds(dec["bits"])
    data = d.audio()
    nib = adpcm.unpack(data, len(data) * 2, order=dec["order"])
    traj = acc_trajectory(nib, dec)

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

    # ---- 1. THE CENSUS.  What a re-PLAY costs at every frame boundary the
    # container has, which is every point a branch in this design can land on:
    # 56.3's targets are frame indices and this player seeks to a frame.
    den = 2 * d.fps
    # THE INDEX IS A NIBBLE INDEX AND THE POSITION IS A BYTE ONE, which is the
    # one conversion in this file and it is worth the line: getting it wrong
    # reads the trajectory at HALF the target and produces a census that is
    # entirely plausible -- a distribution of the right shape over the wrong
    # instants.  The cross-check against the machine's own branch point below is
    # what caught it.
    pos = lambda f: 2 * (f * d.aud_hz // den) - 1
    frames = [f for f in range(1, d.nframes) if pos(f) < len(traj)]
    dcs = [abs(dec["init"] - traj[pos(f)][0]) for f in frames]
    idxs = [traj[pos(f)][1] for f in frames]
    dcs_s = sorted(dcs)
    p = lambda q: dcs_s[min(len(dcs_s) - 1, int(q * len(dcs_s)))]
    print(f"--- 1. A RE-PLAYED BRANCH COSTS `init - acc(target)`, EXACTLY AND "
          f"FOREVER.  {len(frames)} frame boundaries of {a.container}:")
    print(f"     |DC| against the {hi}-unit clamp: mean {statistics.mean(dcs):.1f} "
          f"({statistics.mean(dcs)*100/hi:.1f}%), median {statistics.median(dcs):.0f}, "
          f"p90 {p(0.90)}, worst {max(dcs)} ({max(dcs)*100/hi:.1f}%)")
    print(f"     ...and the step index the encoder assumed at those points runs "
          f"{min(idxs)}..{max(idxs)} of 48, median {statistics.median(idxs):.0f} "
          f"-- a re-PLAY sets it to 0, so a branch into a LOUD passage gets the "
          f"offset AND a step index that has to climb back")
    # The machine run's own branch, so the two layers are checked against each
    # other rather than merely agreeing in prose.
    F37 = 37
    dc37 = dec["init"] - traj[pos(F37)][0]
    print(f"     frame {F37}, the branch tools/bench/packed_run.sh runs on the "
          f"machine: DC {dc37} -- and MAME's capture measured the chip at "
          f"exactly that, constant over 62,500 samples (FINDINGS 71.3)")
    ck(abs(dc37) == 65,
       f"the host's arithmetic for the machine's own branch point is {abs(dc37)} "
       f"and the capture said 65 -- one number, two layers")

    # ---- 2. THE DECAY THAT ISN'T.  A re-PLAY's error is constant BY
    # CONSTRUCTION -- same step index, same nibbles, one offset -- and playing
    # through is not, because the step indices differ too.  The point of
    # measuring it here is that the constancy is a PROPERTY OF THE PREDICTOR
    # and not of the ten seconds this container happens to hold.
    print(f"--- 2. AND IT DOES NOT DECAY.  The accumulator is an integrator with "
          f"no leak: a re-PLAY changes the STARTING value and nothing else, so "
          f"the same nibbles produce the same deltas and the offset is carried "
          f"to the end of the stream. The machine agrees -- AC 0.00 over four "
          f"seconds (FINDINGS 71.3). Playing THROUGH the branch is worse and is "
          f"not constant, because the step index differs as well: -355 falling "
          f"to -108 over four seconds, which is clamping and not forgetting.")

    # ---- 3. THE ENCODER'S FIX, PRICED.  Reset the predictor where a branch can
    # land and a re-PLAYing player is EXACT rather than close.
    if not os.path.exists(a.raw):
        print(f"     (no {a.raw}: the encoder trade below needs the source PCM)")
        return 1 if fails else 0
    import struct
    pcm = struct.unpack("<%dh" % (os.path.getsize(a.raw) // 2),
                        open(a.raw, "rb").read())
    src12 = [max(-2048, min(2047, x >> 4)) for x in pcm][:len(nib)]
    per_frame = d.aud_hz // den * 2            # samples in one frame slot
    print(f"--- 3. THE ONLY FIX THAT MAKES A BRANCH FREE IS THE ENCODER'S, and "
          f"here is its bill.  Reset the predictor every N frames when encoding; "
          f"a player that re-PLAYs at a branch landing on one of those points is "
          f"then EXACT, not close:")
    print(f"       {"reset every":>24}  {'SNR dB':>8}  {'vs shipped':>10}  "
          f"{'branch points made free':>24}")
    base = None
    rows = []
    for label, period in [("never (shipped)", None),
                          (f"{d.cad_f} frames (the cadence)", d.cad_f * per_frame),
                          ("1 frame", per_frame)]:
        nb = encode_reset(src12, dec, period)
        got = decode_reset(nb, dec, period)
        v = snr(src12, got)
        if base is None:
            base = v
        free = (0 if period is None
                else (len(frames) // d.cad_f if period != per_frame
                      else len(frames)))
        rows.append((label, v, v - base, free))
        print(f"       {label:>24}  {v:8.2f}  {v-base:+10.2f}  "
              f"{free:>15} of {len(frames)}")
    # THE ASSERTION IS THE ORDER AND THE SIGN, not the decibel: the source PCM
    # is a property of the disc and the encoder is greedy, so the exact figures
    # move with the window. What must not move is that resetting COSTS SNR and
    # that resetting more often costs more -- if it ever came out free, the
    # predictor would not be doing anything and the codec would be pointless.
    ck(rows[1][1] <= rows[0][1] + 1e-9 and rows[2][1] <= rows[1][1] + 1e-9,
       f"resetting the predictor costs SNR, and resetting it more often costs "
       f"more: {rows[0][1]:.2f} -> {rows[1][1]:.2f} -> {rows[2][1]:.2f} dB")
    ck(rows[2][1] > rows[0][1] - 3.0,
       f"...and a reset EVERY FRAME is {rows[0][1]-rows[2][1]:.2f} dB, which is "
       f"the price of making all {len(frames)} of this container's branch points "
       f"exact. The step table's floor is a constant 16 and the recursion "
       f"re-converges in a few samples, which is why twelve resets a second is "
       f"not twelve times anything")

    print(f"--- 4. WHAT THIS DOES NOT SETTLE.")
    print(f"     * Nothing here is a rate and nothing here ran on silicon. The "
          f"branch behaviour is MAME's okim6258 -- PLAY sets the accumulator to "
          f"-2, the step index to 0 and the nibble select to 0 -- which is the "
          f"model FINDINGS 66 fitted to the machine and NOT a measurement of an "
          f"MSM6258V. It joins session 34's fifth hardware item.")
    print(f"     * The census is ONE container, ten seconds, one passage at "
          f"-13.4 dBFS (FINDINGS 69). The offset a re-PLAY costs is the signal's "
          f"own value at the cut, so a louder passage costs more, up to the "
          f"clamp -- and the disc peaks at 946 of 2048 (69.2).")
    print(f"     * The reset-every-frame encode is NOT in tools/encoder. It is "
          f"a container change (a DLXP3), it costs bytes nothing and SNR "
          f"something, and what decides it is whether a branch is allowed to "
          f"land anywhere or only on frames the encoder was told about.")
    print("AUDIO SEEK GATE " + ("GREEN" if not fails else f"RED: {len(fails)}"))
    return 1 if fails else 0


if __name__ == "__main__":
    sys.exit(main())