#!/usr/bin/env python3
"""Encode a clip from REAL captured footage with HALF-PEL integer motion
compensation, and measure what half-pel buys over whole-pel.

Real footage has *sub-pixel* motion: a camera pan rarely moves an exact integer
number of pixels per frame. Whole-pixel block MC can only ever snap to the
nearest integer shift, so the leftover half-pixel shows up as residual energy
(blur ghosts at edges). Half-pel MC predicts at the half-grid by integer
bilinear interpolation — `(a+b+1)>>1` on an axis, `(a+b+c+d+2)>>2` on the
diagonal — which is the EXACT integer op `wai::int_video::predict()` mirrors on
the decode side, so the clip stays byte-exact (frame-state-equivalence) while the
residual (and the bitstream) shrinks.

The "camera": a sub-pixel Ken-Burns pan over a real photograph (bilinear sampled
— float is fine here, this is capture; only the DECODE is integer). We then encode
the SAME clip twice — whole-pel (mv_precision=1) and half-pel (mv_precision=2) —
and print the rANS byte delta. The half-pel bitstream + its per-frame/clip hashes
are written out; the Rust decoder reproduces them bit-for-bit from the committed
`mv_precision` field.

  ffmpeg -i "<a real photo>.heic" /tmp/real_full.png      # capture source
  python tools/wai_int_video_real_encode.py /tmp/real_full.png <out-dir> [prefix] [size] [nf]
"""
import sys, os, json, base64
import numpy as np
from compressai.ans import RansEncoder

BLOCK = 16       # motion block size
SEARCH = 4       # whole-pixel MV search range (±), in WHOLE pixels
Q = 5            # residual quantization step (fine enough to preserve the sub-pixel gain)
RMAX = (255 + Q - 1) // Q


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


def fit_cdf(syms):
    """Build a precision-16 CDF over residual symbols [-RMAX, RMAX] FIT to the
    actual residual histogram (with a floor so every symbol stays decodable, plus
    the escape slot). The bitstream carries it, so the decoder uses the identical
    model — a data-fit CDF just makes the byte count reflect true residual entropy
    instead of a mismatched constant, so the whole-vs-half comparison is honest."""
    n = 2 * RMAX + 1
    hist = np.bincount(np.asarray(syms, np.int64) + RMAX, minlength=n).astype(np.float64)
    total = 1 << 16
    f = np.maximum((hist / hist.sum() * (total - n)).astype(np.int64), 1)  # floor 1 each
    f[f.argmax()] += (total - 1) - int(f.sum())                            # reals sum to total-1
    cdf = [0]
    for v in f:
        cdf.append(cdf[-1] + int(v))
    cdf.append(total)                                                      # escape slot
    return cdf


def bilinear_sample(src, oy, ox, h, w):
    """Sample an h×w window from float src at sub-pixel top-left (oy,ox) — the
    CAMERA. Bilinear, rounded to u8. (Float lives only on the capture side.)"""
    H, W, _ = src.shape
    ys = oy + np.arange(h)
    xs = ox + np.arange(w)
    y0 = np.floor(ys).astype(np.int64); x0 = np.floor(xs).astype(np.int64)
    wy = (ys - y0)[:, None, None]; wx = (xs - x0)[None, :, None]
    y0c = np.clip(y0, 0, H - 1); y1c = np.clip(y0 + 1, 0, H - 1)
    x0c = np.clip(x0, 0, W - 1); x1c = np.clip(x0 + 1, 0, W - 1)
    a = src[np.ix_(y0c, x0c)]; b = src[np.ix_(y0c, x1c)]
    c = src[np.ix_(y1c, x0c)]; d = src[np.ix_(y1c, x1c)]
    top = a * (1 - wx) + b * wx
    bot = c * (1 - wx) + d * wx
    return np.clip(np.round(top * (1 - wy) + bot * wy), 0, 255).astype(np.int64)


def predict_block(ref, by, bx, bh, bw, dy, dx, prec, H, W):
    """The EXACT integer op `int_video::predict()` runs, vectorized over a block.
    `dy,dx` are in whole pels when prec==1, half pels when prec==2. Each tap is
    clamped independently (matching the decoder's per-access clamp)."""
    if prec == 1:
        ys = np.clip(np.arange(by, by + bh) + dy, 0, H - 1)
        xs = np.clip(np.arange(bx, bx + bw) + dx, 0, W - 1)
        return ref[np.ix_(ys, xs)]
    iy = np.arange(by, by + bh) + (dy >> 1)   # arithmetic shift == floor/2 (neg-safe)
    ix = np.arange(bx, bx + bw) + (dx >> 1)
    fy, fx = dy & 1, dx & 1                    # 0/1 even for negatives (two's complement)

    def g(yo, xo):
        return ref[np.ix_(np.clip(iy + yo, 0, H - 1), np.clip(ix + xo, 0, W - 1))]
    if fy == 0 and fx == 0:
        return g(0, 0)
    if fy == 0:
        return (g(0, 0) + g(0, 1) + 1) >> 1
    if fx == 0:
        return (g(0, 0) + g(1, 0) + 1) >> 1
    return (g(0, 0) + g(0, 1) + g(1, 0) + g(1, 1) + 2) >> 2


def search_mvs(cur, ref, prec, H, W):
    """Two-stage block search. Stage 1: whole-pel SAD over [-SEARCH,SEARCH]^2.
    Stage 2 (prec==2 only): refine ±1 half-pel around 2×(best whole-pel). SAD sums
    over the 3 channels (one MV per block). Returns MVs in `prec`-th-pel units."""
    mvs = []
    for by in range(0, H, BLOCK):
        bh = min(BLOCK, H - by)
        for bx in range(0, W, BLOCK):
            bw = min(BLOCK, W - bx)
            cb = cur[by:by + bh, bx:bx + bw]
            best, bmv = None, (0, 0)
            for dy in range(-SEARCH, SEARCH + 1):
                for dx in range(-SEARCH, SEARCH + 1):
                    pred = predict_block(ref, by, bx, bh, bw, dy, dx, 1, H, W)
                    sad = int(np.abs(cb - pred).sum())
                    if best is None or sad < best:
                        best, bmv = sad, (dy, dx)
            if prec == 1:
                mvs.append([bmv[0], bmv[1]])
                continue
            cy, cx = 2 * bmv[0], 2 * bmv[1]   # center in half-pel units
            best2, bhalf = None, (cy, cx)
            for hy in (cy - 1, cy, cy + 1):
                for hx in (cx - 1, cx, cx + 1):
                    pred = predict_block(ref, by, bx, bh, bw, hy, hx, 2, H, W)
                    sad = int(np.abs(cb - pred).sum())
                    if best2 is None or sad < best2:
                        best2, bhalf = sad, (hy, hx)
            mvs.append([bhalf[0], bhalf[1]])
    return mvs


def encode(frames, prec, H, W):
    """Encode the clip at the given MV precision; return (rANS bytes, mvs, hashes,
    cdf, offset). Residuals (hence reconstruction/hashes) are independent of the
    entropy model, so we gather all symbols first, then fit the CDF to them."""
    offset = -RMAX
    syms, mvs_all, recon, hashes = [], [], [], []
    for f in range(len(frames)):
        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], prec, H, W)
            pred = mc_full(recon[-1], mvs, prec, H, W)
        rq = np.round((frames[f] - pred) / Q).astype(np.int64).clip(-RMAX, RMAX)
        syms.extend(rq.reshape(-1).tolist())
        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'))
    cdf = fit_cdf(syms)
    bits = RansEncoder().encode_with_indexes(syms, [0] * len(syms), [cdf], [len(cdf)], [offset])
    return bytes(bits), mvs_all, hashes, cdf, offset


def mc_full(ref, mvs, prec, H, W):
    """Apply per-block MVs over a full 3-channel reference (decoder-mirroring)."""
    nbx = (W + BLOCK - 1) // BLOCK
    pred = np.empty((H, W, 3), np.int64)
    for by in range(0, H, BLOCK):
        bh = min(BLOCK, H - by)
        for bx in range(0, W, BLOCK):
            bw = min(BLOCK, W - bx)
            dy, dx = mvs[(by // BLOCK) * nbx + (bx // BLOCK)]
            for c in range(3):
                pred[by:by + bh, bx:bx + bw, c] = predict_block(ref[..., c], by, bx, bh, bw, dy, dx, prec, H, W)
    return pred


def main(src_png, outdir, prefix='video_real', size=96, nf=8):
    os.makedirs(outdir, exist_ok=True)
    from PIL import Image
    full = np.asarray(Image.open(src_png).convert('RGB')).astype(np.float64)
    Hs, Ws, _ = full.shape
    H = W = size
    # A detailed, textured region (avoid flat sky) + a sub-pixel Ken-Burns pan.
    oy0, ox0 = Hs * 0.55, Ws * 0.55           # start window top-left (real terrain detail)
    vy, vx = 1.3, 0.8                          # px/frame — deliberately NON-integer
    frames = [bilinear_sample(full, oy0 + vy * f, ox0 + vx * f, H, W) for f in range(nf)]

    out = {}
    for prec, tag in ((1, 'whole'), (2, 'half')):
        bits, mvs_all, hashes, cdf, offset = encode(frames, prec, H, W)
        clip_hash = format(fnv1a(b''.join(bytes.fromhex(h) for h in hashes)), '016x')
        out[tag] = (bits, mvs_all, hashes, cdf, offset, clip_hash)

    wb, hb = len(out['whole'][0]), len(out['half'][0])
    print(f'real footage {nf}x {H}x{W} (sub-pixel pan v=({vy},{vx}) px/frame):')
    print(f'  whole-pel rANS: {wb} bytes')
    print(f'  half-pel  rANS: {hb} bytes   ->  {100*(wb-hb)/wb:+.1f}%  ({wb-hb:+d} bytes)')

    # Ship the HALF-PEL clip (smaller, and the point of the exercise).
    bits, mvs_all, hashes, cdf, offset, clip_hash = out['half']
    json.dump({'block': BLOCK, 'q': Q, 'rmax': RMAX, 'offset': offset, 'cdf': cdf,
               'n_frames': nf, 'h': H, 'w': W, 'mv_precision': 2,
               'mvs': mvs_all, 'rans_b64': base64.b64encode(bits).decode()},
              open(f'{outdir}/{prefix}_bitstream.json', 'w'))
    json.dump({'n_frames': nf, 'h': H, 'w': W, 'frame_hashes': hashes, 'clip_hash': clip_hash,
               'whole_pel_bytes': wb, 'half_pel_bytes': hb},  # so the demo can state the real saving
              open(f'{outdir}/{prefix}.json', 'w'))
    print(f'  half-pel clip_hash {clip_hash}  (I-frame {hashes[0][:8]})  -> {prefix}_bitstream.json')


if __name__ == '__main__':
    main(sys.argv[1], sys.argv[2],
         sys.argv[3] if len(sys.argv) > 3 else 'video_real',
         int(sys.argv[4]) if len(sys.argv) > 4 else 96,
         int(sys.argv[5]) if len(sys.argv) > 5 else 8)