#!/usr/bin/env python3
"""Encode a short clip with a deterministic INTEGER inter-frame codec and emit a
bitstream whose decode is byte-exact on every machine — the temporal extension of
"conformance is a hash" (frame-state-equivalence).

The codec: an I-frame predicted from a constant, then P-frames predicted by
integer whole-pixel block motion compensation from the previous *reconstructed*
frame; the residual is integer-quantized and rANS-coded (CompressAI's coder, the
byte-exact match of the Rust `RansDecoder`). Every step of the DECODE — motion
compensation (array indexing), dequant, add, clamp — is integer, so the
reconstructed frames are identical on every machine. (Lossy vs the source, like
int_hyper; byte-exact vs each other.)

  python tools/wai_int_video_encode.py <out-dir>
"""
import sys, os, json, base64
import numpy as np
from compressai.ans import RansEncoder

BLOCK = 8        # motion block size
SEARCH = 5       # whole-pixel MV search range (±)
Q = 6            # residual quantization step
RMAX = (255 + Q - 1) // Q  # residual symbol half-range (covers the worst case, no clamp loss)
NF, H, W = 5, 48, 48       # frames / height / width


def fnv1a(buf):
    h = 14695981039346656037
    for b in buf:
        h = ((h ^ b) * 1099511628211) & 0xFFFFFFFFFFFFFFFF
    return h


def make_clip():
    """A deterministic RGB clip: a bright disc translating over a gradient — clear
    motion for the block MC to exploit (small residuals)."""
    frames = []
    yy, xx = np.mgrid[0:H, 0:W]
    grad = ((xx * 90 // W) + (yy * 70 // H)).astype(np.int64)  # background gradient
    for f in range(NF):
        cx, cy, r = 16 + 6 * f, 18 + 4 * f, 12
        disc = ((xx - cx) ** 2 + (yy - cy) ** 2) <= r * r
        img = np.empty((H, W, 3), np.int64)
        img[..., 0] = np.clip(40 + grad, 0, 255)
        img[..., 1] = np.clip(60 + grad // 2, 0, 255)
        img[..., 2] = np.clip(90 + grad // 3, 0, 255)
        img[disc] = [250, 210, 70]  # the disc
        frames.append(img.astype(np.int64))
    return frames


def mc_predict(ref, mvs):
    """Build a prediction by shifting each BLOCK of `ref` by its (dy,dx), with
    clamped source coordinates — the exact integer op the Rust decoder mirrors."""
    pred = np.empty_like(ref)
    nbx = (W + BLOCK - 1) // BLOCK
    for by in range(0, H, BLOCK):
        for bx in range(0, W, BLOCK):
            dy, dx = mvs[(by // BLOCK) * nbx + (bx // BLOCK)]
            for y in range(by, min(by + BLOCK, H)):
                sy = min(max(y + dy, 0), H - 1)
                for x in range(bx, min(bx + BLOCK, W)):
                    sx = min(max(x + dx, 0), W - 1)
                    pred[y, x] = ref[sy, sx]
    return pred


def search_mvs(cur, ref):
    """Per block, pick the integer MV in [-SEARCH, SEARCH]^2 minimizing SAD against
    the reconstructed reference. (Encoder side; the decoder just applies the MVs.)"""
    mvs = []
    for by in range(0, H, BLOCK):
        for bx in range(0, W, BLOCK):
            cb = cur[by:by + BLOCK, bx:bx + BLOCK]
            best, bmv = None, (0, 0)
            for dy in range(-SEARCH, SEARCH + 1):
                for dx in range(-SEARCH, SEARCH + 1):
                    ys = np.clip(np.arange(by, by + cb.shape[0]) + dy, 0, H - 1)
                    xs = np.clip(np.arange(bx, bx + cb.shape[1]) + dx, 0, W - 1)
                    rb = ref[np.ix_(ys, xs)]
                    sad = int(np.abs(cb - rb).sum())
                    if best is None or sad < best:
                        best, bmv = sad, (dy, dx)
            mvs.append([int(bmv[0]), int(bmv[1])])
    return mvs


def laplacian_cdf():
    """A peaked-at-zero quantized CDF over residual symbols [-RMAX, RMAX] (+escape),
    precision 16; the model constant both sides share."""
    n = 2 * RMAX + 1
    pmf = np.exp(-np.abs(np.arange(n) - RMAX) / 6.0)
    total = 1 << 16
    f = np.maximum((pmf / pmf.sum() * (total - 1)).round().astype(np.int64), 1)
    f[f.argmax()] += (total - 1) - int(f.sum())   # real symbols sum to total-1
    cdf = [0]
    for v in f:
        cdf.append(cdf[-1] + int(v))
    cdf.append(total)                              # escape slot → cdf[n+1] = 2^16
    return cdf


def main(outdir):
    os.makedirs(outdir, exist_ok=True)
    frames = make_clip()
    cdf = laplacian_cdf()
    offset = -RMAX
    enc = RansEncoder()
    syms, mvs_all, recon, hashes = [], [], [], []

    for f in range(NF):
        if f == 0:
            pred = np.full((H, W, 3), 128, np.int64)   # I-frame: predict from grey
            mvs = []
        else:
            mvs = search_mvs(frames[f], recon[-1])
            pred = np.stack([mc_predict(recon[-1][..., c], mvs) for c in range(3)], -1)
        resid = frames[f] - pred
        rq = np.round(resid / Q).astype(np.int64).clip(-RMAX, RMAX)   # quantized residual
        syms.extend(rq.reshape(-1).tolist())                          # raster: y, x, c
        # reconstruct exactly as the decoder will (dequant + add + clamp)
        rec = np.clip(pred + rq * Q, 0, 255).astype(np.int64)
        recon.append(rec)
        mvs_all.append(mvs)
        hashes.append(format(fnv1a(np.ascontiguousarray(rec.astype(np.uint8)).tobytes()), '016x'))

    bits = enc.encode_with_indexes(syms, [0] * len(syms), [cdf], [len(cdf)], [offset])
    clip_hash = format(fnv1a(b''.join(bytes.fromhex(h) for h in hashes)), '016x')

    json.dump({'block': BLOCK, 'q': Q, 'rmax': RMAX, 'offset': offset, 'cdf': cdf,
               'n_frames': NF, 'h': H, 'w': W,
               'mvs': mvs_all, 'rans_b64': base64.b64encode(bytes(bits)).decode()},
              open(f'{outdir}/int_video_bitstream.json', 'w'))
    json.dump({'n_frames': NF, 'h': H, 'w': W, 'frame_hashes': hashes, 'clip_hash': clip_hash},
              open(f'{outdir}/int_video.json', 'w'))
    print(f'int_video: {NF}x {H}x{W} RGB, {len(bits)} rANS bytes + {sum(len(m) for m in mvs_all)} MVs '
          f'-> clip_hash {clip_hash}')
    print('per-frame:', ' '.join(h[:8] for h in hashes))


if __name__ == '__main__':
    main(sys.argv[1])