"""The registered watermark schemes `wai-mark:tent64` (extensions/synthesis.md §7.5) and
`wai-mark:tent64s` (§7.9), and the attacks of their battery (§7.5.7, §7.9.6), written from the
extension's text.

Shared by the two Python verifiers of the synthesis corpora (synthesis_verify.py and
synthesis_mark_verify.py); it shares no code with wai-rs. Standard library only.

Frames are `bytes` of RGB8, frame after frame. Every function is integer arithmetic, so its
results are the extension's to the byte.
"""

import hashlib
import math

SCHEME = "wai-mark:tent64"
GRID = 64
CELLS = GRID * GRID
LENGTH = 248          # the shortened code's length (§7.5.3)
PARITY = 184
T = 29                # the errors the code corrects
CELLS_PER_BIT = 16
AMPLITUDE = 12        # §7.5.4
MARGIN = (3, 2)       # 2·σ·T ≥ 3·Q
MIN_SIDE = 128        # §7.5.1
MAX_FRAME_BYTES = 1 << 28
MAX_INT = (1 << 53) - 1
CLASS_MAX = 3


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


def stream(label, seed=None):
    """§7.5.3: 32-bit big-endian words, eight from each SHA-256(label ‖ [seed ‖] n)."""
    n = 0
    while True:
        d = hashlib.sha256(label + (b"" if seed is None else seed.to_bytes(4, "big")) + n.to_bytes(4, "big")).digest()
        for o in range(0, 32, 4):
            yield int.from_bytes(d[o:o + 4], "big")
        n += 1


def _layout():
    r = stream(b"wai-mark:tent64 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]
    signs = [1 if next(r) % 2 == 0 else -1 for _ in range(CELLS)]
    mask = [next(r) & 1 for _ in range(LENGTH)]
    return perm, signs, mask


PERM, SIGNS, MASK = _layout()
BIT_CELLS = [PERM[CELLS_PER_BIT * j: CELLS_PER_BIT * j + CELLS_PER_BIT] for j in range(LENGTH)]

# ---------------------------------------------------------------- the code, §7.5.3

EXP = [0] * 510
LOG = [0] * 256
_x = 1
for _i in range(255):
    EXP[_i] = _x
    LOG[_x] = _i
    _x <<= 1
    if _x & 0x100:
        _x ^= 0x11D        # x^8 + x^4 + x^3 + x^2 + 1
for _i in range(255, 510):
    EXP[_i] = EXP[_i - 255]


def gmul(a, b):
    return 0 if a == 0 or b == 0 else EXP[LOG[a] + LOG[b]]


def generator():
    """The least common multiple of the minimal polynomials of α^1 … α^58, as an integer whose
    bit d is the coefficient of x^d."""
    g, cosets = 1, []
    for i in range(1, 2 * T + 1):
        coset, j = [], i % 255
        while j not in coset:
            coset.append(j)
            j = 2 * j % 255
        if sorted(coset) in cosets:
            continue
        cosets.append(sorted(coset))
        p = [1]                                   # coefficients from x^0 up
        for j in coset:
            q = [0] * (len(p) + 1)
            for d, c in enumerate(p):
                q[d + 1] ^= c
                q[d] ^= gmul(c, EXP[j])
            p = q
        assert all(c in (0, 1) for c in p), "a minimal polynomial is binary"
        pm = sum(c << d for d, c in enumerate(p))
        prod, a, b = 0, g, pm                     # product over GF(2)
        while b:
            if b & 1:
                prod ^= a
            a <<= 1
            b >>= 1
        g = prod
    return g


G = generator()


def encode(payload):
    """The codeword of an 8-byte payload: bit j the coefficient of x^(247 − j)."""
    m = int.from_bytes(payload, "big")
    r = m << PARITY
    for d in range(r.bit_length() - 1, PARITY - 1, -1):
        if r >> d & 1:
            r ^= G << (d - PARITY)
    c = (m << PARITY) | r
    return [c >> (LENGTH - 1 - j) & 1 for j in range(LENGTH)]


def _syndromes(c):
    return [_eval(c, i) for i in range(1, 2 * T + 1)]


def _eval(c, i):
    v = 0
    e = 0
    while c:
        if c & 1:
            v ^= EXP[i * e % 255]
        c >>= 1
        e += 1
    return v


def decode(word):
    """Bounded-distance decoding within 29: (payload, bits corrected), or None where no codeword
    is within 29 bits or the payload's first byte is above 3."""
    c = 0
    for b in word:
        c = (c << 1) | b
    s = _syndromes(c)
    errors = 0
    if any(s):
        # The error locator: the shortest linear recurrence the syndromes satisfy, over GF(2^8).
        lam, prev, L, m, bd = [1], [1], 0, 1, 1
        for n in range(2 * T):
            d = s[n]
            for i in range(1, min(L, len(lam) - 1) + 1):
                d ^= gmul(lam[i], s[n - i])
            if d == 0:
                m += 1
                continue
            coef = gmul(d, EXP[(255 - LOG[bd]) % 255])
            old = lam[:]
            lam = lam + [0] * max(0, len(prev) + m - len(lam))
            for i, bv in enumerate(prev):
                lam[i + m] ^= gmul(coef, bv)
            if 2 * L <= n:
                L, prev, bd, m = n + 1 - L, old, d, 1
            else:
                m += 1
        while len(lam) > 1 and lam[-1] == 0:
            lam.pop()
        if len(lam) - 1 != L or L > T:
            return None
        roots = []
        for e in range(255):
            v = 0
            for i, co in enumerate(lam):
                if co:
                    v ^= gmul(co, EXP[(-e * i) % 255])
            if v == 0:
                roots.append(e)
        if len(roots) != L or any(e >= LENGTH for e in roots):
            return None
        for e in roots:
            c ^= 1 << e
        if any(_syndromes(c)):
            return None
        errors = L
    payload = (c >> PARITY).to_bytes(8, "big")
    if payload[0] > CLASS_MAX:
        return None
    return payload, errors


# ---------------------------------------------------------------- frames, §7.5.1 and §7.5.2

def admits(w, h, n):
    return w >= MIN_SIDE and h >= MIN_SIDE and n >= 1 and w * h * 3 <= MAX_FRAME_BYTES and w * h * 3 * n <= MAX_INT


def luma(r, g, b):
    return (77 * r + 150 * g + 29 * b + 128) >> 8


def tent(u, length):
    c = 2 * u + 1 - length
    sg = (c > 0) - (c < 0)
    return sg * (length - abs(2 * abs(c) - length))


def axis(n):
    """Per sample: (cell index, tent, cell length); and the edges."""
    edges = [i * n // GRID for i in range(GRID + 1)]
    cell, tents, lens = [], [], []
    for i in range(GRID):
        L = edges[i + 1] - edges[i]
        for x in range(edges[i], edges[i + 1]):
            cell.append(i)
            tents.append(tent(x - edges[i], L))
            lens.append(L)
    return cell, tents, lens, edges


def luma_plane(frame):
    return [luma(frame[p], frame[p + 1], frame[p + 2]) for p in range(0, len(frame), 3)]


def bit_statistics(y, w, h, layout=None):
    """T_j over one luma plane of w × h, under a layout (bit cells, signs) or tent64's."""
    bit_cells, signs = (BIT_CELLS, SIGNS) if layout is None else (layout[0], layout[1])
    cx, fx, _, _ = axis(w)
    cy, fy, _, _ = axis(h)
    s = [0] * CELLS
    for row in range(h):
        acc = [0] * GRID
        base = row * w
        for col in range(w):
            acc[cx[col]] += fx[col] * y[base + col]
        r0, f = cy[row] * GRID, fy[row]
        for i in range(GRID):
            s[r0 + i] += f * acc[i]
    return [sum(signs[k] * s[k] for k in bit_cells[j]) for j in range(LENGTH)]


# ---------------------------------------------------------------- embedding, §7.5.4

def embed(unmarked, w, h, n, payload, layout=None):
    """The marked frames, and each frame's amplitudes, under a layout (bit cells, signs, mask)
    or tent64's."""
    bit_cells, signs, mask = (BIT_CELLS, SIGNS, MASK) if layout is None else layout
    code = encode(payload)
    sigma = [1 if code[j] ^ mask[j] else -1 for j in range(LENGTH)]
    _, fx, lx, ex = axis(w)
    _, fy, ly, ey = axis(h)
    fb = w * h * 3
    out = bytearray(unmarked)
    amps = []
    for f in range(n):
        frame = unmarked[f * fb:(f + 1) * fb]
        chosen = []
        for j in range(LENGTH):
            sg = sigma[j]
            pixels = []          # (offset, ε·ω, z, ν)
            q = 0
            for k in bit_cells[j]:
                i, jj = k % GRID, k // GRID
                eps = signs[k]
                for row in range(ey[jj], ey[jj + 1]):
                    for col in range(ex[i], ex[i + 1]):
                        om = fx[col] * fy[row]
                        q += abs(om)
                        pixels.append(((row * w + col) * 3, eps * om, sg * eps * om, lx[col] * ly[row]))

            def at(a):
                t, vals = 0, []
                for p, eo, z, nu in pixels:
                    d = (z > 0) - (z < 0)
                    d *= (2 * a * abs(z) + nu) // (2 * nu)
                    r = min(255, max(0, frame[p] + d))
                    g = min(255, max(0, frame[p + 1] + d))
                    b = min(255, max(0, frame[p + 2] + d))
                    t += eo * luma(r, g, b)
                    vals.append((p, r, g, b))
                return t, vals

            pick = None
            for a in range(AMPLITUDE + 1):
                t, vals = at(a)
                if MARGIN[1] * sg * t >= MARGIN[0] * q:
                    pick = (a, vals)
                    break
            if pick is None:
                pick = (AMPLITUDE, at(AMPLITUDE)[1])
            a, vals = pick
            for p, r, g, b in vals:
                out[f * fb + p] = r
                out[f * fb + p + 1] = g
                out[f * fb + p + 2] = b
            chosen.append(a)
        amps.append(chosen)
    return bytes(out), amps


# ---------------------------------------------------------------- detection, §7.5.5

def orient(y, w, h, o):
    ow, oh = (h, w) if o & 4 else (w, h)
    out = [0] * (w * h)
    for yy in range(oh):
        y1 = oh - 1 - yy if o & 2 else yy
        for xx in range(ow):
            x1 = ow - 1 - xx if o & 1 else xx
            sx, sy = (y1, x1) if o & 4 else (x1, y1)
            out[yy * ow + xx] = y[sy * w + sx]
    return out, ow, oh


def read_word(samples, w, h, n, o, layout=None):
    mask = MASK if layout is None else layout[2]
    fb = w * h * 3
    t = [0] * LENGTH
    for f in range(n):
        y, ow, oh = orient(luma_plane(samples[f * fb:(f + 1) * fb]), w, h, o)
        for j, v in enumerate(bit_statistics(y, ow, oh, layout)):
            t[j] += v
    return [(1 if t[j] > 0 else 0) ^ mask[j] for j in range(LENGTH)]


def detect(samples, w, h, n):
    """(payload, orientation, errors) of the first orientation that decodes, or None."""
    if not admits(w, h, n) or len(samples) != w * h * 3 * n:
        return None
    for o in range(8):
        got = decode(read_word(samples, w, h, n, o))
        if got is not None:
            return got[0], o, got[1]
    return None


# ---------------------------------------------------------------- the battery, §7.5.7

def rnd(n, d):
    """n / d rounded to the nearest integer, halves away from zero."""
    sg = (n > 0) - (n < 0)
    return sg * ((2 * abs(n) + d) // (2 * d))


def clamp(v):
    return 0 if v < 0 else 255 if v > 255 else v


def int_motion(frame, q):
    return bytes(clamp(128 + rnd(v - 128, q) * q) for v in frame)


def upscale2(px, w, h):
    """wai-synthesis:upscale2, §4.1."""
    out = bytearray(4 * w * h * 3)
    ow = 2 * w
    for y in range(h):
        y1 = min(y + 1, h - 1)
        for x in range(w):
            x1 = min(x + 1, w - 1)
            for c in range(3):
                a, b = px[(y * w + x) * 3 + c], px[(y * w + x1) * 3 + c]
                cc, d = px[(y1 * w + x) * 3 + c], px[(y1 * w + x1) * 3 + c]
                out[((2 * y) * ow + 2 * x) * 3 + c] = a
                out[((2 * y) * ow + 2 * x + 1) * 3 + c] = (a + b + 1) >> 1
                out[((2 * y + 1) * ow + 2 * x) * 3 + c] = (a + cc + 1) >> 1
                out[((2 * y + 1) * ow + 2 * x + 1) * 3 + c] = (a + b + cc + d + 2) >> 2
    return bytes(out)


def denoise3(px, w, h):
    """wai-synthesis:denoise3, §4.2."""
    k = ((1, 2, 1), (2, 4, 2), (1, 2, 1))
    out = bytearray(w * h * 3)
    for y in range(h):
        for x in range(w):
            for c in range(3):
                s = 8
                for j in range(3):
                    yy = min(max(y + j - 1, 0), h - 1)
                    for i in range(3):
                        xx = min(max(x + i - 1, 0), w - 1)
                        s += k[j][i] * px[(yy * w + xx) * 3 + c]
                out[(y * w + x) * 3 + c] = s >> 4
    return bytes(out)


def _weights(n, m):
    out = []
    for d in range(m):
        lo, hi = d * n, (d + 1) * n
        row = []
        for s in range(lo // m, -(-hi // m)):
            o = min(hi, (s + 1) * m) - max(lo, s * m)
            if o > 0:
                row.append((s, o))
        out.append(row)
    return out


def area(px, w, h, nw, nh):
    wx, wy = _weights(w, nw), _weights(h, nh)
    den = w * h
    out = bytearray(nw * nh * 3)
    for dy, ry in enumerate(wy):
        for dx, rx in enumerate(wx):
            for c in range(3):
                num = 0
                for sy, oy in ry:
                    for sx, ox in rx:
                        num += oy * ox * px[(sy * w + sx) * 3 + c]
                out[(dy * nw + dx) * 3 + c] = (2 * num + den) // (2 * den)
    return bytes(out)


ROTATIONS = {"rotate-0.5": (65534, 572), "rotate-1": (65526, 1144), "rotate-2": (65496, 2287), "rotate-5": (65287, 5712)}

# §7.9.6: the simulated captures of a screen, (num, den) of the resampling and the rotation.
SCREENS = {"screen-a": ((5, 4), (65526, 1144)), "screen-b": ((4, 5), (65496, -2287))}


def rotate(px, w, h, cs, sn):
    one = 1 << 16

    def p(x, y, c):
        return 0 if x < 0 or y < 0 or x >= w or y >= h else px[(y * w + x) * 3 + c]

    out = bytearray(w * h * 3)
    for y in range(h):
        for x in range(w):
            u, v = 2 * x + 1 - w, 2 * y + 1 - h
            X = (cs * u + sn * v + (w - 1) * one) >> 1
            Y = (-sn * u + cs * v + (h - 1) * one) >> 1
            xi, fx, yi, fy = X >> 16, X & 0xFFFF, Y >> 16, Y & 0xFFFF
            for c in range(3):
                val = (p(xi, yi, c) * (one - fx) * (one - fy) + p(xi + 1, yi, c) * fx * (one - fy)
                       + p(xi, yi + 1, c) * (one - fx) * fy + p(xi + 1, yi + 1, c) * fx * fy)
                out[(y * w + x) * 3 + c] = (val + (1 << 31)) >> 32
    return bytes(out)


def attack(name, px, w, h, seed):
    """One picture attack of §7.5.7 on a single frame: (result, width, height)."""
    if name == "none":
        return px, w, h
    if name.startswith("int-motion-q"):
        return int_motion(px, int(name[len("int-motion-q"):])), w, h
    if name == "staged-base-q8":
        hw, hh = w // 2, h // 2
        half = bytearray(hw * hh * 3)
        for y in range(hh):
            for x in range(hw):
                for c in range(3):
                    s = sum(px[((2 * y + dy) * w + 2 * x + dx) * 3 + c] for dy in (0, 1) for dx in (0, 1))
                    half[(y * hw + x) * 3 + c] = (s + 2) >> 2
        return upscale2(int_motion(bytes(half), 8), hw, hh), 2 * hw, 2 * hh
    if name == "denoise3":
        return denoise3(px, w, h), w, h
    if name == "denoise3-x4":
        for _ in range(4):
            px = denoise3(px, w, h)
        return px, w, h
    if name.startswith("scale-"):
        n, d = map(int, name[len("scale-"):].split("/"))
        nw, nh = w * n // d, h * n // d
        return area(px, w, h, nw, nh), nw, nh
    if name == "upscale2":
        return upscale2(px, w, h), 2 * w, 2 * h
    if name == "aspect-3/4":
        nw = w * 3 // 4
        return area(px, w, h, nw, h), nw, h
    if name == "half-then-int-motion-q16":
        nw, nh = w // 2, h // 2
        return int_motion(area(px, w, h, nw, nh), 16), nw, nh
    for side, prefix in enumerate(("crop-left-", "crop-right-", "crop-top-", "crop-bottom-")):
        if name.startswith(prefix) and name.endswith("%"):
            pc = int(name[len(prefix):-1])
            cx, cy = w * pc // 100, h * pc // 100
            x0, y0, nw, nh = [(cx, 0, w - cx, h), (0, 0, w - cx, h), (0, cy, w, h - cy), (0, 0, w, h - cy)][side]
            out = b"".join(px[(y * w + x0) * 3:(y * w + x0 + nw) * 3] for y in range(y0, y0 + nh))
            return out, nw, nh
    if name in SCREENS:
        (n, d), (cs, sn) = SCREENS[name]
        nw, nh = w * n // d, h * n // d
        shown = rotate(area(px, w, h, nw, nh), nw, nh, cs, sn)
        shown = bytes((2 * v * (v + 255) + 510) // 1020 for v in shown)
        r = stream(b"wai-mark:battery noise", seed)
        return bytes(clamp(v + next(r) % 9 - 4) for v in shown), nw, nh
    if name == "periodic":
        qw, qh = w // 2, h // 2
        return b"".join(px[((y % qh) * w + x % qw) * 3:((y % qh) * w + x % qw) * 3 + 3] for y in range(h) for x in range(w)), w, h
    if name.startswith("crop-"):
        pc = int(name[len("crop-"):-1])
        dx, dy = w * pc // 200, h * pc // 200
        nw, nh = w - 2 * dx, h - 2 * dy
        out = b"".join(px[((y * w) + dx) * 3:((y * w) + dx + nw) * 3] for y in range(dy, dy + nh))
        return out, nw, nh
    if name == "mirror":
        return b"".join(px[(y * w + x) * 3:(y * w + x) * 3 + 3] for y in range(h) for x in range(w - 1, -1, -1)), w, h
    if name == "rotate-90":
        return b"".join(px[((h - 1 - xp) * w + yp) * 3:((h - 1 - xp) * w + yp) * 3 + 3] for yp in range(w) for xp in range(h)), h, w
    if name in ROTATIONS:
        cs, sn = ROTATIONS[name]
        return rotate(px, w, h, cs, sn), w, h
    if name.startswith("noise-"):
        s = int(name[len("noise-"):])
        r = stream(b"wai-mark:battery noise", seed)
        return bytes(clamp(v + next(r) % (2 * s + 1) - s) for v in px), w, h
    if name.startswith("brightness"):
        b = int(name[len("brightness"):])
        return bytes(clamp(v + b) for v in px), w, h
    if name.startswith("contrast-"):
        n, d = map(int, name[len("contrast-"):].split("/"))
        return bytes(clamp(128 + rnd((v - 128) * n, d)) for v in px), w, h
    if name == "local-blur":
        out = bytearray(px)
        x0, x1, y0, y1 = w // 4, 3 * w // 4, h // 4, 3 * h // 4
        for by in range(y0, y1, 16):
            for bx in range(x0, x1, 16):
                ex, ey = min(bx + 16, x1), min(by + 16, y1)
                cnt = (ex - bx) * (ey - by)
                for c in range(3):
                    tot = sum(px[(y * w + x) * 3 + c] for y in range(by, ey) for x in range(bx, ex))
                    mval = (2 * tot + cnt) // (2 * cnt)
                    for y in range(by, ey):
                        for x in range(bx, ex):
                            out[(y * w + x) * 3 + c] = mval
        return bytes(out), w, h
    raise ValueError("no picture attack named %r" % name)


def attack_frames(name, px, w, h, n, seed):
    """A geometric attack, or `periodic`, on each of n frames: (result, width, height)."""
    fb = w * h * 3
    outs = [attack(name, px[f * fb:(f + 1) * fb], w, h, seed) for f in range(n)]
    return b"".join(o[0] for o in outs), outs[0][1], outs[0][2]


# ================================================================ wai-mark:tent64s, §7.9

SCHEME_S = "wai-mark:tent64s"
TILE = 32
POSITIONS = TILE * TILE
PER_TILE = 4
STAGE2_FRAMES = 8
WORKING_SIDE = 1024
ROBUST_CLAMP = 100
CANDIDATES = 8
DECODES = 4
SPAN = ((7, 20), (3, 4))
SLOPE = (3, 20)
SUB = 64


def _tile_layout():
    """§7.9.2: the shuffle of the tile's 1024 positions, their signs and the mask."""
    r = stream(b"wai-mark:tent64s stream")
    perm = list(range(POSITIONS))
    for i in range(POSITIONS - 1, 0, -1):
        j = next(r) % (i + 1)
        perm[i], perm[j] = perm[j], perm[i]
    signs = [1 if next(r) % 2 == 0 else -1 for _ in range(POSITIONS)]
    mask = [next(r) & 1 for _ in range(LENGTH)]
    return perm, signs, mask


T_PERM, T_SIGNS, T_MASK = _tile_layout()


def _cell_of(p, t):
    """The frame cell of tile position p in tile t (left to right, then top to bottom)."""
    row, col = divmod(p, TILE)
    return GRID * (TILE * (t // 2) + row) + TILE * (t % 2) + col


_S_SIGNS = [0] * CELLS
for _p in range(POSITIONS):
    for _t in range(4):
        _S_SIGNS[_cell_of(_p, _t)] = T_SIGNS[_p]
LAYOUT_S = ([[_cell_of(T_PERM[PER_TILE * j + k], t) for k in range(PER_TILE) for t in range(4)] for j in range(LENGTH)], _S_SIGNS, T_MASK)


def embed_s(unmarked, w, h, n, payload):
    """§7.9.3: §7.5.4's embedding under the tiled layout."""
    return embed(unmarked, w, h, n, payload, LAYOUT_S)


def _rnd(n, d):
    """n / d to the nearest integer, halves away from zero, for any d ≠ 0."""
    if d < 0:
        n, d = -n, -d
    sg = (n > 0) - (n < 0)
    return sg * ((2 * abs(n) + d) // (2 * d))


def _working(frame, w, h, k):
    y = luma_plane(frame)
    if k == 1:
        return y, w, h
    ww, hh, kk = w // k, h // k, k * k
    out = []
    for oy in range(hh):
        for ox in range(ww):
            s = sum(y[(oy * k + dy) * w + ox * k + dx] for dy in range(k) for dx in range(k))
            out.append((s + kk // 2) // kk)
    return out, ww, hh


def _robust(y, w, h):
    row = [0] * (w * h)
    for r in range(h):
        b = r * w
        for x in range(w):
            row[b + x] = sum(y[b + min(max(x + d, 0), w - 1)] for d in (-2, -1, 0, 1, 2))
    out = [0] * (w * h)
    for r in range(h):
        rows = [min(max(r + d, 0), h - 1) * w for d in (-2, -1, 0, 1, 2)]
        for x in range(w):
            v = 25 * y[r * w + x] - sum(row[q + x] for q in rows)
            out[r * w + x] = -ROBUST_CLAMP if v < -ROBUST_CLAMP else ROBUST_CLAMP if v > ROBUST_CLAMP else v
    return out


def _sign_plane(r, w, h):
    lx, ly = max(2, w // GRID), max(2, h // GRID)
    wq, hq = w - lx + 1, h - ly + 1
    fx = [tent(u, lx) for u in range(lx)]
    fy = [tent(u, ly) for u in range(ly)]
    horiz = [sum(fx[u] * r[row * w + x + u] for u in range(lx)) for row in range(h) for x in range(wq)]
    q = []
    for y in range(hq):
        for x in range(wq):
            s = sum(fy[v] * horiz[(y + v) * wq + x] for v in range(ly))
            q.append((s > 0) - (s < 0))
    return q, wq, hq


def _autocorrelation(planes, wq, hq):
    """§7.9.4 step 4, exactly: one product of two integers per frame (rows laid end to end at
    a stride of 2·Wq, four bytes a value), summed, read back. Returns A(dx, dy) as a function."""
    import sys as _sys
    stride = 2 * wq
    n = hq * stride
    ones = int.from_bytes(b"\x01\x00\x00\x00" * n, "little")
    total = 0
    for q in planes:
        flat = [0] * n
        for y in range(hq):
            flat[y * stride:y * stride + wq] = q[y * wq:(y + 1) * wq]
        fwd = b"".join((v + 1).to_bytes(4, "little") for v in flat)
        rev = b"".join((v + 1).to_bytes(4, "little") for v in reversed(flat))
        a = int.from_bytes(fwd, "little") - ones
        b = int.from_bytes(rev, "little") - ones
        total += a * b
    bias = int.from_bytes(b"\x00\x00\x00\x80" * (2 * n - 1), "little")
    raw = (total + bias).to_bytes(4 * (2 * n - 1), "little")
    import array
    slots = array.array("I")
    slots.frombytes(raw)
    if _sys.byteorder == "big":
        slots.byteswap()

    def at(dx, dy):
        if abs(dx) >= wq or abs(dy) >= hq:
            return 0
        return slots[n - 1 + dy * stride + dx] - (1 << 31)
    return at


def _candidates(at, wq, hq):
    """§7.9.4 step 5: up to 8 lattice candidates (dx, dy, e, g)."""
    (ln, ld), (hn, hd) = SPAN

    def ov(x, y):
        return 0 if abs(x) >= wq or abs(y) >= hq else (wq - abs(x)) * (hq - abs(y))
    rows = []
    idx = 0
    for dx in range(-(-ln * wq // ld), hn * wq // hd + 1):
        m = SLOPE[0] * dx // SLOPE[1]
        a1, o1 = None, None
        for dy in range(-m, m + 1):
            a1, o1 = at(dx, dy), ov(dx, dy)
            for e in range(-(-ln * hq // ld), hn * hq // hd + 1):
                g = _rnd(-dy * e, dx)
                a = a1 + at(g, e) + at(dx + g, dy + e) + at(dx - g, dy - e)
                o = o1 + ov(g, e) + ov(dx + g, dy + e) + ov(dx - g, dy - e)
                rows.append((-((65536 * a) // max(1, math.isqrt(o))), idx, dx, dy, e, g))
                idx += 1
    rows.sort()
    out = []
    for _, _, dx, dy, e, g in rows:
        if all(max(abs(t[0] - dx), abs(t[1] - dy), abs(t[2] - e)) > 2 for t in out):
            out.append((dx, dy, e, g))
            if len(out) == CANDIDATES:
                break
    return out


def _vertex(l, c, r):
    d = l - 2 * c + r
    if d >= 0:
        return 0
    return max(-32, min(32, _rnd(32 * (r - l), -d)))


def _refine(at, dx, dy):
    c = at(dx, dy)
    return SUB * dx + _vertex(at(dx - 1, dy), c, at(dx + 1, dy)), SUB * dy + _vertex(at(dx, dy - 1), c, at(dx, dy + 1))


def _tile_stats(f, mx, my, px, py):
    nx, ny = TILE * mx, TILE * my
    tx = [tent(u, mx) for u in range(mx)]
    ty = [tent(u, my) for u in range(my)]
    rowsum = []                                       # per bin row, per cell column
    for row in range(ny):
        base = row * nx
        rowsum.append([sum(tx[u] * f[base + (a * mx + px + u) % nx] for u in range(mx)) for a in range(TILE)])
    s = [0] * POSITIONS
    for b in range(TILE):
        rs = [rowsum[(b * my + py + v) % ny] for v in range(my)]
        for a in range(TILE):
            s[b * TILE + a] = sum(ty[v] * rs[v][a] for v in range(my))
    return s


def _oriented(s, o):
    sg = -1 if ((o & 1) ^ ((o >> 1) & 1)) else 1
    out = [0] * POSITIONS
    for y in range(TILE):
        for x in range(TILE):
            x1 = TILE - 1 - x if o & 1 else x
            y1 = TILE - 1 - y if o & 2 else y
            x2, y2 = (y1, x1) if o & 4 else (x1, y1)
            out[y * TILE + x] = sg * s[y2 * TILE + x2]
    return out


_FULL = (1 << POSITIONS) - 1
_ROW_LO = [sum(((1 << (TILE - c)) - 1) << (TILE * r) for r in range(TILE)) for c in range(TILE)]
_ROW_HI = [_FULL ^ _ROW_LO[c] for c in range(TILE)]


def _roll(bm, r, c):
    """The bitmap over offsets s = 32·y0 + x0 whose bit s is bm's bit at position
    32·((r + y0) mod 32) + (c + x0) mod 32."""
    k = (TILE * r) % POSITIONS
    v = ((bm >> k) | (bm << (POSITIONS - k))) & _FULL
    if c:
        v = ((v >> c) & _ROW_LO[c]) | ((v << (TILE - c)) & _ROW_HI[c])
    return v


def _add_plane(counter, plane):
    i = 0
    while plane:
        if i == len(counter):
            counter.append(0)
        counter[i], plane = counter[i] ^ plane, counter[i] & plane
        i += 1


def _agreements(so):
    """For every offset, the agreement count of §7.9.4 step 8: over every bit, the pairs of its
    four cells whose ε·S are both above 0 or both not. Bit-sliced over the 1024 offsets."""
    pos = sum(1 << p for p in range(POSITIONS) if so[p] > 0)
    neg = sum(1 << p for p in range(POSITIONS) if so[p] < 0)
    odd_c, eq_c = [], []
    for j in range(LENGTH):
        planes = []
        for k in range(PER_TILE):
            p = T_PERM[PER_TILE * j + k]
            planes.append(_roll(pos if T_SIGNS[p] > 0 else neg, p // TILE, p % TILE))
        a, b, c, d = planes
        _add_plane(odd_c, a ^ b ^ c ^ d)
        _add_plane(eq_c, (a & b & c & d) | (_FULL & ~(a | b | c | d)))
    out = []
    for s in range(POSITIONS):
        odd = sum(((odd_c[i] >> s) & 1) << i for i in range(len(odd_c)))
        eq = sum(((eq_c[i] >> s) & 1) << i for i in range(len(eq_c)))
        out.append(2 * LENGTH + odd + 4 * eq)
    return out


def _word_s(so, x0, y0):
    w = []
    for j in range(LENGTH):
        t = 0
        for k in range(PER_TILE):
            p = T_PERM[PER_TILE * j + k]
            t += T_SIGNS[p] * so[((p // TILE + y0) % TILE) * TILE + (p % TILE + x0) % TILE]
        w.append((1 if t > 0 else 0) ^ T_MASK[j])
    return w


def stage2(samples, w, h, n, words_only=False, trace=False):
    """§7.9.4 stage 2: (payload, orientation, errors, resync) or None; with words_only, the
    words it decodes, all of them; with trace, every candidate kept (its vectors, lattice, bins
    and phase) and the hypotheses decoded, each [agreement, tile, orientation, x0, y0]."""
    k = -(-max(w, h) // WORKING_SIDE)
    fb = w * h * 3
    robs, planes = [], []
    for f in range(min(n, STAGE2_FRAMES)):
        y, ww, hh = _working(samples[f * fb:(f + 1) * fb], w, h, k)
        r = _robust(y, ww, hh)
        q, wq, hq = _sign_plane(r, ww, hh)
        robs.append(r)
        planes.append(q)
    at = _autocorrelation(planes, wq, hq)
    tiles = []
    for ci, (dx, dy, e, g) in enumerate(_candidates(at, wq, hq)):
        x1, y1 = _refine(at, dx, dy)
        x2, y2 = _refine(at, g, e)
        det = x1 * y2 - x2 * y1
        if det <= 0:
            continue
        mx, my = (min(8, max(4, v // (SUB * TILE))) for v in (x1, y2))
        nx, ny = TILE * mx, TILE * my
        sx, sy = SUB * nx, SUB * ny
        folds = [[0] * (nx * ny) for _ in range(4)]
        for r in robs:
            for yy in range(hh):
                for xx in range(ww):
                    u = sx * (y2 * xx - x2 * yy) // det
                    v = sy * (x1 * yy - y1 * xx) // det
                    folds[(u // nx) % 2 + 2 * ((v // ny) % 2)][(v % ny) * nx + u % nx] += r[yy * ww + xx]
        best = None
        for py in range(my):
            for px in range(mx):
                parts = [_tile_stats(fc, mx, my, px, py) for fc in folds]
                s = [parts[0][i] + parts[1][i] + parts[2][i] + parts[3][i] for i in range(POSITIONS)]
                en = sum(s[i] * s[i] - sum(p[i] * p[i] for p in parts) for i in range(POSITIONS))
                if best is None or en > best[0]:
                    best = (en, (px, py), s)
        tiles.append(({"lattice": [x1, y1, x2, y2], "bins": [mx, my], "phase": list(best[1]), "candidate": ci}, best[2], [dx, dy, g, e]))
    hyps = []
    oriented = {}
    for ti, (_, s, _v) in enumerate(tiles):
        for o in range(8):
            so = _oriented(s, o)
            oriented[(ti, o)] = so
            for i, cnt in enumerate(_agreements(so)):
                hyps.append((-cnt, len(hyps), ti, o, i % TILE, i // TILE))
    hyps.sort()
    if trace:
        return {"candidates": [{"vectors": v, "lattice": r["lattice"], "bins": r["bins"], "phase": r["phase"]} for r, _, v in tiles],
                "decoded": [[-h[0], h[2], h[3], h[4], h[5]] for h in hyps[:DECODES]]}
    words = []
    for rank, (ncnt, _, ti, o, x0, y0) in enumerate(hyps[:DECODES]):
        wd = _word_s(oriented[(ti, o)], x0, y0)
        words.append(wd)
        if words_only:
            continue
        got = decode(wd)
        if got is not None:
            res = dict(tiles[ti][0], offset=[x0, y0], rank=rank, agreement=-ncnt)
            return got[0], o, got[1], res
    return words if words_only else None


def stage2_trace(samples, w, h, n):
    """Stage 2's trace (§7.9.4 steps 5 to 9), whatever stage 1 reads; None for frames the
    scheme does not admit."""
    if not admits(w, h, n) or len(samples) != w * h * 3 * n:
        return None
    return stage2(samples, w, h, n, trace=True)


def detect_s(samples, w, h, n):
    """§7.9.4: (payload, stage, orientation, errors, resync) or None."""
    if not admits(w, h, n) or len(samples) != w * h * 3 * n:
        return None
    for o in range(8):
        got = decode(read_word(samples, w, h, n, o, LAYOUT_S))
        if got is not None:
            return got[0], 1, o, got[1], None
    got = stage2(samples, w, h, n)
    if got is None:
        return None
    return got[0], 2, got[1], got[2], got[3]