open-standards/wai/live-conformance/_wai_audio_mark.py

"""wai-mark:band64, the audio scheme of extensions/synthesis.md §7.10, written from that text.

A second implementation: it shares no code with wai-rs. The code (the BCH code of §7.5.3) is
_wai_mark.py's, which is written from §7.5. Standard library only.
"""

import hashlib
import math
import operator

import _wai_mark as TENT

SCHEME = "wai-mark:band64"
RATES = (8000, 12000, 16000, 22050, 24000, 32000, 44100, 48000)
MAX_CHANNELS = 8
MAX_LENGTH = 1 << 28
BANDS = (10, 16, 24, 34, 48, 66, 89, 120, 160)
NB = 8
PAIRS = 4
SLOTS = 62
CELLS = PAIRS * SLOTS
PERIOD = 4 * SLOTS
STEP = 848
HALF = STEP // 2
OFFSET = STEP // 4
GATE = 3400
GAIN_MAX = 320
PASSES = 3
PHASES = 4
ONE = 16384
LENGTH = TENT.LENGTH
LABEL = b"wai-mark:band64 stream"


def sha256(b):
    return hashlib.sha256(b).digest()


def hop(rate):
    return rate // 50


def admits(rate, channels, length):
    return rate in RATES and 1 <= channels <= MAX_CHANNELS and rate <= length <= MAX_LENGTH


# ---------------------------------------------------------------- the cosine tables (§7.10.1)

_TABLES = {}


def table(n):
    """rnd(16384·cos(2πi / 4n)), i from 0 to 4n − 1."""
    t = _TABLES.get(n)
    if t is None:
        t = [int(math.floor(ONE * math.cos(2 * math.pi * i / (4 * n)) + 0.5)) for i in range(4 * n)]
        _TABLES[n] = t
    return t


def tables_bytes():
    out = bytearray()
    for r in RATES:
        for v in table(2 * hop(r)):
            out += (v & 0xFFFF).to_bytes(2, "big")
    return bytes(out)


class Basis:
    def __init__(self, rate):
        h = hop(rate)
        n = 2 * h
        t = table(n)
        m = 4 * n
        self.n, self.hop = n, h
        self.window = [t[(n - 2 * x - 1) % m] for x in range(n)]
        self.cos, self.sin, self.wc, self.ws = [], [], [], []
        for k in range(BANDS[0], BANDS[NB]):
            c = [t[(4 * k * x) % m] for x in range(n)]
            s = [t[(4 * k * x - n) % m] for x in range(n)]
            self.cos.append(c)
            self.sin.append(s)
            self.wc.append([a * b for a, b in zip(self.window, c)])
            self.ws.append([a * b for a, b in zip(self.window, s)])


_BASES = {}


def basis(rate):
    if rate not in _BASES:
        _BASES[rate] = Basis(rate)
    return _BASES[rate]


# ---------------------------------------------------------------- the layout (§7.10.3)

def _stream():
    n = 0
    while True:
        d = sha256(LABEL + n.to_bytes(4, "big"))
        for o in range(8):
            yield int.from_bytes(d[4 * o:4 * o + 4], "big")
        n += 1


def _layout():
    r = _stream()
    perm = list(range(CELLS))
    for i in range(CELLS - 1, 0, -1):
        j = next(r) % (i + 1)
        perm[i], perm[j] = perm[j], perm[i]
    sign = [1 if next(r) % 2 == 0 else -1 for _ in range(CELLS)]
    mask = [next(r) & 1 for _ in range(LENGTH)]
    bit = [0] * CELLS
    for j, q in enumerate(perm):
        bit[q] = j
    return perm, bit, sign, mask


PERM, BIT, SIGN, MASK = _layout()


# ---------------------------------------------------------------- samples (§7.10.1, §7.10.2)

def unpack(buf, channels):
    """The canonical buffer: S16LE, channels interleaved, as one list per channel."""
    vals = [int.from_bytes(buf[i:i + 2], "little", signed=True) for i in range(0, len(buf), 2)]
    return [vals[c::channels] for c in range(channels)]


def pack(chans):
    out = bytearray()
    for vs in zip(*chans):
        for v in vs:
            out += (v & 0xFFFF).to_bytes(2, "little")
    return bytes(out)


def mono(chans):
    c = len(chans)
    if c == 1:
        return list(chans[0])
    return [sum(vs) // c for vs in zip(*chans)]


def lg(e):
    """⌊256·log2 e⌋ by eight squarings of the 62-bit mantissa; 0 for 0."""
    if e == 0:
        return 0
    k = e.bit_length() - 1
    m = e >> (k - 62) if k >= 62 else e << (62 - k)
    frac = 0
    for _ in range(8):
        m = (m * m) >> 62
        frac <<= 1
        if m >= 1 << 63:
            frac |= 1
            m >>= 1
    return 256 * k + frac


def shift(a):
    return 2 * (lg((1024 + a) ** 2) - lg((1024 - a) ** 2))


_SHIFTS = [shift(a) for a in range(GAIN_MAX + 1)]


def gain_for(want):
    need = abs(want)
    a = GAIN_MAX
    for i, s in enumerate(_SHIFTS):
        if s >= need:
            a = i
            break
    return a if want > 0 else -a if want < 0 else 0


def frame_count(rate, length, start):
    n = 2 * hop(rate)
    return 0 if length < start + n else (length - start - n) // hop(rate) + 1


def band_energies(rate, x, start):
    b = basis(rate)
    out = []
    mul = operator.mul
    for f in range(frame_count(rate, len(x), start)):
        s = start + f * b.hop
        fr = x[s:s + b.n]
        e = [0] * NB
        for band in range(NB):
            acc = 0
            for k in range(BANDS[band] - BANDS[0], BANDS[band + 1] - BANDS[0]):
                re = sum(map(mul, fr, b.wc[k]))
                im = sum(map(mul, fr, b.ws[k]))
                acc += re * re + im * im
            e[band] = acc
        out.append(e)
    return out


def gate(e):
    """§7.10.2: the median log of every tile, every band at every frame but the last, less 3400."""
    logs = sorted(lg(e[f][b] + e[f + 1][b]) for f in range(len(e) - 1) for b in range(NB))
    return logs[len(logs) // 2] - GATE if logs else 0


def cell(e, f, i, g=None):
    a0 = e[f][2 * i] + e[f + 1][2 * i]
    b0 = e[f + 2][2 * i] + e[f + 3][2 * i]
    a1 = e[f][2 * i + 1] + e[f + 1][2 * i + 1]
    b1 = e[f + 2][2 * i + 1] + e[f + 3][2 * i + 1]
    if not (a0 and b0 and a1 and b1):
        return None
    la0, lb0, la1, lb1 = lg(a0), lg(b0), lg(a1), lg(b1)
    if g is not None and min(la0, lb0, la1, lb1) < g:
        return None
    return (la0 - lb0) - (la1 - lb1)


def soft(q, d):
    phi = (SIGN[q] * d - OFFSET) % STEP
    return min(phi, STEP - phi) - abs(phi - HALF)


def rnd(n, d):
    s = 1 if n > 0 else -1 if n < 0 else 0
    return s * ((2 * abs(n) + d) // (2 * d))


# ---------------------------------------------------------------- embedding (§7.10.4)

def synthesize(rate, chans, cells):
    """chans with each cell's tiles scaled: cells are (frame, pair, gain)."""
    b = basis(rate)
    length = len(chans[0])
    nf = frame_count(rate, length, 0)
    gain = [[0] * NB for _ in range(nf)]
    for f, i, a in cells:
        for g in range(4):
            s = a if g < 2 else -a
            gain[f + g][2 * i] = s
            gain[f + g][2 * i + 1] = -s
    out = []
    mul = operator.mul
    den = b.n << 49
    for x in chans:
        m = [0] * length
        for f in range(nf):
            g = gain[f]
            if not any(g):
                continue
            s0 = f * b.hop
            fr = x[s0:s0 + b.n]
            v = [0] * b.n
            for band in range(NB):
                if g[band] == 0:
                    continue
                for k in range(BANDS[band] - BANDS[0], BANDS[band + 1] - BANDS[0]):
                    gre = g[band] * sum(map(mul, fr, b.wc[k]))
                    gim = g[band] * sum(map(mul, fr, b.ws[k]))
                    ck, sk = b.cos[k], b.sin[k]
                    v = [vt + gre * c + gim * s for vt, c, s in zip(v, ck, sk)]
            for t in range(b.n):
                m[s0 + t] += rnd(b.window[t] * v[t], den)
        out.append([max(-32768, min(32767, xv + rnd(mv, 1 << 16))) for xv, mv in zip(x, m)])
    return out


def embed(rate, chans, payload):
    """The marked channels of chans under payload, and each carrying cell's last gain."""
    code = TENT.encode(payload)
    e = band_energies(rate, mono(chans), 0)
    g = gate(e)
    cells = []
    for u in range(len(e) // 4):
        for i in range(PAIRS):
            q = SLOTS * i + u % SLOTS
            d = cell(e, 4 * u, i, g)
            if d is None:
                continue
            t = code[BIT[q]] ^ MASK[BIT[q]]
            tau = SIGN[q] * d
            o = OFFSET + t * HALF
            z = o + STEP * ((tau - o + HALF) // STEP)
            cells.append([4 * u, i, q, z, z - tau])
    gains = [0] * len(cells)
    y = chans
    for p in range(PASSES):
        gains = [gain_for(SIGN[c[2]] * c[4]) for c in cells]
        y = synthesize(rate, chans, [(c[0], c[1], g) for c, g in zip(cells, gains)])
        if p + 1 < PASSES:
            e2 = band_energies(rate, mono(y), 0)
            for c in cells:
                d2 = cell(e2, c[0], c[1])
                err = c[3] - SIGN[c[2]] * (d2 if d2 is not None else 0)
                c[4] = max(-STEP, min(STEP, c[4] + rnd(4 * err, 5)))
    return y, gains


# ---------------------------------------------------------------- detection (§7.10.5)

def phase_start(rate, p):
    return p * hop(rate) // PHASES


def words(rate, chans):
    """Every alignment's word, in the detector's order: (offset, phase, word)."""
    x = mono(chans)
    stats = [None] * PHASES

    def phase(p):
        if stats[p] is None:
            e = band_energies(rate, x, phase_start(rate, p))
            g = gate(e)
            per = []
            for r in range(4):
                occ = []
                s = r
                while s + 3 < len(e):
                    occ.append([cell(e, s, i, g) for i in range(PAIRS)])
                    s += 4
                per.append(occ)
            stats[p] = per
        return stats[p]

    for o in range(PERIOD):
        r = (4 - o % 4) % 4
        k0 = (r + o) // 4
        for p in range(PHASES):
            t = [0] * LENGTH
            for u, cs in enumerate(phase(p)[r]):
                for i, d in enumerate(cs):
                    if d is not None:
                        q = SLOTS * i + (k0 + u) % SLOTS
                        t[BIT[q]] += soft(q, d)
            yield o, p, [(1 if t[j] > 0 else 0) ^ MASK[j] for j in range(LENGTH)]


def detect(rate, channels, buf):
    """(payload, offset, phase, errors) or None."""
    length = len(buf) // (2 * channels) if channels else 0
    if not admits(rate, channels, length) or len(buf) != 2 * channels * length:
        return None
    for o, p, w in words(rate, unpack(buf, channels)):
        d = TENT.decode(w)
        if d is not None:
            return d[0], o, p, d[1]
    return None


def embed_bytes(rate, channels, buf, payload):
    y, gains = embed(rate, unpack(buf, channels), payload)
    return pack(y), gains


# ---------------------------------------------------------------- the battery's audio transforms

def _noise_stream(seed):
    n = 0
    while True:
        d = sha256(b"wai-mark:battery noise" + seed.to_bytes(4, "big") + n.to_bytes(4, "big"))
        for o in range(8):
            yield int.from_bytes(d[4 * o:4 * o + 4], "big")
        n += 1


def clamp16(v):
    return -32768 if v < -32768 else 32767 if v > 32767 else v


def _ulaw_encode(v):
    bias, clip = 0x84, 32635
    sign = 0x80 if v < 0 else 0
    v = min(-v if v < 0 else v, clip) + bias
    exp = 7
    mask = 0x4000
    while exp > 0 and not (v & mask):
        exp -= 1
        mask >>= 1
    mant = (v >> (exp + 3)) & 0x0F
    return ~(sign | (exp << 4) | mant) & 0xFF


def _ulaw_decode(u):
    u = ~u & 0xFF
    t = (((u & 0x0F) << 3) + 0x84) << ((u & 0x70) >> 4)
    return (0x84 - t) if (u & 0x80) else (t - 0x84)


def area(x, n_from, n_to):
    """Area resampling of one channel from n_from samples a second to n_to."""
    length = len(x) * n_to // n_from
    out = []
    for X in range(length):
        lo, hi = X * n_from, (X + 1) * n_from          # in units of 1/(n_from·n_to) s, scaled by n_to
        acc = 0
        s = lo // n_to
        while s * n_to < hi:
            ov = min(hi, (s + 1) * n_to) - max(lo, s * n_to)
            if ov > 0 and s < len(x):
                acc += ov * x[s]
            s += 1
        out.append(clamp16(rnd(acc, n_from)))
    return out


def attack(name, params, rate, chans, seed):
    """One audio transform of robustness-evidence §4 (the pinned ones): (rate, channels)."""
    if name == "none":
        return rate, [list(c) for c in chans]
    if name == "audio-gain":
        n, d = params
        return rate, [[clamp16(rnd(v * n, d)) for v in c] for c in chans]
    if name == "audio-trim":
        a, b = params
        s, e = rate * a // 1000, rate * b // 1000
        return rate, [c[s:len(c) - e] for c in chans]
    if name == "audio-resample":
        (r,) = params
        return r, [area(c, rate, r) for c in chans]
    if name == "audio-resample-trip":
        (r,) = params
        return rate, [area(area(c, rate, r), r, rate) for c in chans]
    if name == "audio-noise":
        (s,) = params
        w = _noise_stream(seed)
        out = [list(c) for c in chans]
        for t in range(len(chans[0])):
            for c in range(len(chans)):
                out[c][t] = clamp16(out[c][t] + next(w) % (2 * s + 1) - s)
        return rate, out
    if name == "audio-mix-noise":
        n, d = params
        w = _noise_stream(seed)
        noise = [[0] * len(chans[0]) for _ in chans]
        for t in range(len(chans[0])):
            for c in range(len(chans)):
                noise[c][t] = next(w) % 65536 - 32768
        px = sum(v * v for c in chans for v in c)
        pw = sum(v * v for c in noise for v in c)
        if pw == 0:
            return rate, [list(c) for c in chans]
        k = math.isqrt(px * n * (1 << 32) // (pw * d))
        return rate, [[clamp16(v + rnd(u * k, 1 << 16)) for v, u in zip(c, nc)] for c, nc in zip(chans, noise)]
    if name == "audio-mulaw":
        return rate, [[_ulaw_decode(_ulaw_encode(v)) for v in c] for c in chans]
    if name == "audio-requantize":
        (b,) = params
        sh = 16 - b
        return rate, [[(v >> sh) << sh for v in c] for c in chans]
    if name == "audio-echo":
        ms, n, d = params
        lag = rate * ms // 1000
        return rate, [[clamp16(v + (rnd(c[t - lag] * n, d) if t >= lag else 0)) for t, v in enumerate(c)] for c in chans]
    if name == "audio-mono":
        return rate, [mono(chans)]
    raise ValueError("not a pinned audio transform: " + name)