"""The registered watermark scheme `wai-mark:tent64` of extensions/synthesis.md §7.5, and the
attacks of its battery (§7.5.7), 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

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):
    """T_j over one luma plane of w × h."""
    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):
    """The marked frames, and each frame's amplitudes (a string of 248 digits as base-36)."""
    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):
    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)):
            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)}


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
    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)