#!/usr/bin/env python3
"""Write, or check, the vectors of the neural audio route: wai-rs/tests/vectors/neural_audio/.

One vector for each capability whose payload SPEC §5 pins (*Payloads of the neural audio
codecs*): `wai.neural.encodec24`, `wai.neural.encodec48`, `wai.neural.dac16` and
`wai.neural.dac24`. Each directory holds

  decoder.onnx   a tiny decoder with the model's geometry and its decoder contract: the model's
                 codebook count Q, codebooks of 1024, a hop of 320, its channels; inputs
                 audio_codes int64 [1, 1, Q, T] and audio_scales float [1], output audio_values
                 float [1, C, T*320 - d] (d = 8 for DAC 16 and 24 kHz, whose decoders give 8
                 samples fewer, else 0). Its network is not the model's: a code c of codebook i adds
                 c*A[i] + B[i] to the frame's latent z (a row holding 1024 adds nothing), and
                 sample k of the frame on channel ch is tanh(z * W[ch][k]) * scale, the last d
                 samples dropped. A, B and W
                 are dyadic, so every value before the tanh is exact in float32.
  vector.wai     a WAI1 envelope of the capability, pinned to decoder.onnx, whose payload is
                 seeded codes in the pinned layout (some codebooks only, so the route must mark
                 the others unused; for encodec48, two chunks with their own scales).
  expected.f32   the decode at the samples meta.json selects, every channel of each, as
                 little-endian float32, computed here in Python from the payload and the
                 decoder's parameters: the second implementation of the route, which reads the
                 payload, fills the codebooks, runs the network and joins the chunks by the
                 model's linear overlap-add on its own.
  meta.json      what the vector is, and which samples expected.f32 holds: index i where
                 i % every == 0 or i lies in an overlap region.

The ONNX files are written here, protobuf field by field, so the generator needs nothing but
the standard library and its bytes are the same on every machine.

    python3 -I tools/wai_neural_audio_vectors.py           # write the vectors
    python3 -I tools/wai_neural_audio_vectors.py --check   # rebuild them and compare

--check exits 0 when every committed file is what the generator writes (expected.f32 within
1e-6, since it holds a libm tanh), 1 otherwise.
"""
import hashlib
import json
import math
import os
import struct
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "..", "wai-rs", "tests", "vectors", "neural_audio")
K = 1024  # entries per codebook; the value that marks a codebook not used
HOP = 320

# name: capability, sample rate, channels, model codebooks Q, codebooks carried q, frames of each
# chunk, samples per channel n, chunk scales (None: no chunks), chunk stride, the samples the
# decoder gives fewer than frames*320 (DAC 16 and 24 kHz: 8), seed
VECTORS = {
    "encodec24": ("wai.neural.encodec24", 24000, 1, 32, 8, [5], 1523, None, 0, 0, 1),
    "encodec48": ("wai.neural.encodec48", 48000, 2, 16, 4, [150, 6], 49340, [0.75, 0.5], 47520, 0, 4),
    "dac16": ("wai.neural.dac16", 16000, 1, 12, 5, [4], 1201, None, 0, 8, 2),
    "dac24": ("wai.neural.dac24", 24000, 1, 32, 32, [3], 950, None, 0, 8, 3),
}


# ---- the decoder's parameters -------------------------------------------------------------

def params(Q, C):
    """A[i] = a/2^16, B[i] = b/2^8, W[ch][k] = w/2^4: |z| < 4 in steps of 2^-16, and z*W
    takes at most 22 significant bits, so every value before the tanh is exact in float32."""
    A = [[1, -2, 3, -4, 4, -3, 2, -1][i % 8] / 65536.0 for i in range(Q)]
    B = [(((i * 5) % 17) - 8) / 256.0 for i in range(Q)]
    W = [[(((k * 3 + ch * 5) % 31) - 15) / 16.0 for k in range(HOP)] for ch in range(C)]
    return A, B, W


# ---- protobuf, as much of onnx.proto as the decoder needs ---------------------------------

def varint(v):
    out = bytearray()
    v &= (1 << 64) - 1
    while True:
        b = v & 0x7F
        v >>= 7
        if v:
            out.append(b | 0x80)
        else:
            out.append(b)
            return bytes(out)


def fv(field, v):
    return varint(field << 3) + varint(v)


def fb(field, b):
    if isinstance(b, str):
        b = b.encode()
    return varint((field << 3) | 2) + varint(len(b)) + b


FLOAT, INT64, BOOL = 1, 7, 9


def tensor(name, dtype, dims, raw):
    # TensorProto: dims = 1, data_type = 2, name = 8, raw_data = 9.
    return b"".join(fv(1, d) for d in dims) + fv(2, dtype) + fb(8, name) + fb(9, raw)


def f32s(xs):
    return b"".join(struct.pack("<f", x) for x in xs)


def value_info(name, dtype, dims):
    # ValueInfoProto: name = 1, type = 2; TypeProto.tensor_type = 1 {elem_type = 1, shape = 2};
    # TensorShapeProto.dim = 1 {dim_value = 1 | dim_param = 2}.
    shape = b"".join(fb(1, fb(2, d) if isinstance(d, str) else fv(1, d)) for d in dims)
    return fb(1, name) + fb(2, fb(1, fv(1, dtype) + fb(2, shape)))


def attr_int(name, v):
    # AttributeProto: name = 1, i = 3, type = 20 (INT = 2).
    return fb(1, name) + fv(3, v) + fv(20, 2)


def attr_ints(name, vs):
    # ints = 8 (INTS = 7).
    return fb(1, name) + b"".join(fv(8, v) for v in vs) + fv(20, 7)


def node(op, inputs, outputs, attrs=()):
    # NodeProto: input = 1, output = 2, name = 3, op_type = 4, attribute = 5.
    return (b"".join(fb(1, i) for i in inputs) + b"".join(fb(2, o) for o in outputs)
            + fb(3, outputs[0]) + fb(4, op) + b"".join(fb(5, a) for a in attrs))


def decoder_onnx(Q, C, shrink):
    A, B, W = params(Q, C)
    nodes = [
        node("Less", ["audio_codes", "K"], ["present"]),
        node("Cast", ["present"], ["mask"], [attr_int("to", FLOAT)]),
        node("Cast", ["audio_codes"], ["codes_f"], [attr_int("to", FLOAT)]),
        node("Mul", ["codes_f", "A"], ["scaled"]),
        node("Add", ["scaled", "B"], ["shifted"]),
        node("Mul", ["shifted", "mask"], ["used"]),
        node("ReduceSum", ["used", "axes"], ["z"], [attr_int("keepdims", 0)]),
        node("ConvTranspose", ["z", "W"], ["y"], [attr_ints("strides", [HOP]), attr_ints("pads", [0, shrink])]),
        node("Tanh", ["y"], ["bounded"]),
        node("Mul", ["bounded", "audio_scales"], ["audio_values"]),
    ]
    inits = [
        tensor("K", INT64, [], struct.pack("<q", K)),
        tensor("A", FLOAT, [1, 1, Q, 1], f32s(A)),
        tensor("B", FLOAT, [1, 1, Q, 1], f32s(B)),
        tensor("axes", INT64, [1], struct.pack("<q", 2)),
        tensor("W", FLOAT, [1, C, HOP], f32s([w for row in W for w in row])),
    ]
    # GraphProto: node = 1, name = 2, initializer = 5, input = 11, output = 12.
    graph = (b"".join(fb(1, n) for n in nodes) + fb(2, "neural_audio_test_decoder")
             + b"".join(fb(5, t) for t in inits)
             + fb(11, value_info("audio_codes", INT64, [1, 1, Q, "T"]))
             + fb(11, value_info("audio_scales", FLOAT, [1]))
             + fb(12, value_info("audio_values", FLOAT, [1, C, "samples"])))
    # ModelProto: ir_version = 1, producer_name = 2, graph = 7, opset_import = 8 {version = 2}.
    return fv(1, 8) + fb(2, "wai_neural_audio_vectors") + fb(7, graph) + fb(8, fv(2, 17))


# ---- the payload ----------------------------------------------------------------------------

def splitmix64(seed):
    s = seed & ((1 << 64) - 1)
    while True:
        s = (s + 0x9E3779B97F4A7C15) & ((1 << 64) - 1)
        z = s
        z = ((z ^ (z >> 30)) * 0xBF58476D1CE4E5B9) & ((1 << 64) - 1)
        z = ((z ^ (z >> 27)) * 0x94D049BB133111EB) & ((1 << 64) - 1)
        yield z ^ (z >> 31)


def chunks_of(q, frames, seed):
    """Each chunk's codes: q rows of its frame count, codebook-major."""
    g = splitmix64(seed)
    return [[[next(g) % K for _ in range(t)] for _ in range(q)] for t in frames]


def payload(q, frames, n, scales, chunks):
    if scales is None:
        head = struct.pack("<III", q, frames[0], n)
    else:
        head = struct.pack("<IIII", q, frames[-1], n, len(frames)) + f32s(scales)
    return head + b"".join(struct.pack("<H", c) for ch in chunks for row in ch for c in row)


def envelope(cap, kind, sr, C, n, body, digest):
    man = {
        "wai": "1.0", "media": "audio", "intent": "create",
        "model_requirement": {"capability": cap},
        "conditioning": {"kind": kind, "prior": {
            "id": f"neural-audio test decoder {kind}", "sha256": digest,
            "determinism": "none: float inference through ONNX Runtime", "tier": "none"}},
        "target": {"sr": sr, "channels": C, "n_samples": n},
    }
    mb = json.dumps(man, separators=(",", ":")).encode()
    return b"WAI1" + struct.pack("<I", len(mb)) + mb + struct.pack("<I", len(body)) + body


# ---- the second implementation of the route -------------------------------------------------

def parse(cap, Q, C, body, chunked):
    """Read the payload as SPEC §5 lays it out, independently of the generator's writer."""
    if chunked:
        q, t_last, n, count = struct.unpack_from("<IIII", body, 0)
        scales = list(struct.unpack_from(f"<{count}f", body, 16))
        frames = [150] * (count - 1) + [t_last]
        at = 16 + 4 * count
    else:
        q, t, n = struct.unpack_from("<III", body, 0)
        scales, frames, at = [1.0], [t], 12
    chunks = []
    for t in frames:
        rows = []
        for _ in range(q):
            rows.append(list(struct.unpack_from(f"<{t}H", body, at)))
            at += 2 * t
        chunks.append(rows)
    assert at == len(body), f"{cap}: {len(body) - at} bytes after the codes"
    return q, n, scales, chunks


def decode_chunk(Q, C, rows, scale, shrink):
    """The tiny decoder on one chunk: rows past the payload's hold K and add nothing, and the
    last `shrink` samples are not given."""
    A, B, W = params(Q, C)
    t = len(rows[0])
    full = rows + [[K] * t for _ in range(Q - len(rows))]
    z = [0.0] * t
    for i in range(Q):
        for j in range(t):
            c = full[i][j]
            if c < K:
                z[j] += c * A[i] + B[i]
    out = [[math.tanh(z[j] * W[ch][k]) * scale for j in range(t) for k in range(HOP)] for ch in range(C)]
    return [o[:t * HOP - shrink] for o in out]


def overlap_add(decoded, stride):
    """The model's linear overlap-add: a triangle over the first chunk's length."""
    C = len(decoded[0])
    len0 = len(decoded[0][0])
    total = stride * (len(decoded) - 1) + len(decoded[-1][0])
    out = [[0.0] * total for _ in range(C)]
    wsum = [0.0] * total
    for k, ch in enumerate(decoded):
        off = k * stride
        for j in range(len(ch[0])):
            w = 0.5 - abs((j + 1) / (len0 + 1) - 0.5)
            wsum[off + j] += w
            for c in range(C):
                out[c][off + j] += w * ch[c][j]
    return [[v / s for v, s in zip(o, wsum)] for o in out]


def selected(n, every, overlaps):
    return [i for i in range(n) if i % every == 0 or any(a <= i < b for a, b in overlaps)]


def build(name):
    cap, sr, C, Q, q, frames, n, scales, stride, shrink, seed = VECTORS[name]
    onnx = decoder_onnx(Q, C, shrink)
    digest = hashlib.sha256(onnx).hexdigest()
    body = payload(q, frames, n, scales, chunks_of(q, frames, seed))
    wai = envelope(cap, name, sr, C, n, body, digest)
    # The second implementation reads the payload back and decodes it.
    q2, n2, sc, chunks = parse(cap, Q, C, body, scales is not None)
    decoded = [decode_chunk(Q, C, rows, s, shrink) for rows, s in zip(chunks, sc)]
    joined = overlap_add(decoded, stride) if scales is not None else decoded[0]
    every = 16 if scales is not None else 1
    overlaps = [[k * stride, (k - 1) * stride + frames[0] * HOP] for k in range(1, len(frames))]
    at = selected(n2, every, overlaps)
    expected = [joined[c][i] for i in at for c in range(C)]
    meta = {
        "capability": cap, "sample_rate": sr, "channels": C, "model_codebooks": Q,
        "codebooks": q2, "frames": frames, "n_samples": n2, "scales": scales, "shrink": shrink,
        "every": every, "overlaps": overlaps, "expected_values": len(expected),
        "decoder_sha256": digest, "seed": seed,
    }
    return {
        "decoder.onnx": onnx,
        "vector.wai": wai,
        "expected.f32": f32s(expected),
        "meta.json": (json.dumps(meta, indent=2) + "\n").encode(),
    }


def main(argv):
    check = argv[1:] == ["--check"]
    if argv[1:] and not check:
        print(__doc__)
        return 2
    bad = 0
    for name in sorted(VECTORS):
        d = os.path.join(OUT, name)
        for fname, want in build(name).items():
            path = os.path.join(d, fname)
            if not check:
                os.makedirs(d, exist_ok=True)
                with open(path, "wb") as fh:
                    fh.write(want)
                continue
            try:
                with open(path, "rb") as fh:
                    got = fh.read()
            except OSError as e:
                print(f"{name}/{fname}: {e}")
                bad += 1
                continue
            if fname == "expected.f32":
                g = struct.unpack(f"<{len(got) // 4}f", got)
                w = struct.unpack(f"<{len(want) // 4}f", want)
                worst = max((abs(a - b) for a, b in zip(g, w)), default=0.0)
                if len(g) != len(w) or worst > 1e-6:
                    print(f"{name}/{fname}: {len(g)} values, max |diff| {worst:.2e}; want {len(w)} within 1e-6")
                    bad += 1
            elif got != want:
                print(f"{name}/{fname}: differs from what the generator writes")
                bad += 1
        print(f"{name}: {'checked' if check else 'written'}")
    return 1 if bad else 0


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