"""Encode an image to a `wai.neural.mbt2018_mean` WAI envelope, decoded by
`wai-rs/src/neural/hyperprior.rs`.

Mean-scale hyperprior: the envelope carries TWO rANS bitstreams — the
factorized hyper-latent `z`, then the Gaussian-conditional main latent `y`.
CompressAI's rANS is the coder `wai-rs/src/neural/rans.rs` ports (precision=16,
bypass=4) and flushes 32-bit words, so both bitstreams are multiples of 4 bytes,
exactly as the decoder requires.

Wire payload (little-endian, must match hyperprior.rs):
  <IIIIHHBBH>  H, W, Lz, Ly, Cz, Cy, Sz_log2, Sy_log2, q
  Lz bytes  z bitstream
  Ly bytes  y bitstream

Container: WAI1 | u32 manifest_len | manifest JSON | u32 payload_len | payload.

Note on `intent`: this is a LOSSY codec, so "replicate" here means *every sink
reconstructs the identical (lossy) image*, not the original. That cross-sink
identity only holds if the decode is integer-exact (see the determinism work in
container.rs); float ONNX synthesis is "exact on this build" — so we DON'T pin a
prior, and `decode_envelope_strict` will (correctly) flag a replicate claim as
unsound until an integer-exact model is registered. Lenient decode works today.

Run (needs torch + compressai + Pillow):
  python tools/wai_mbt2018_mean_encode.py --quality 3 --size 256 \
    --image in.png --out kodim.mbt2018_mean.wai
"""
from __future__ import annotations

import argparse
import json
import math
import struct
from pathlib import Path

import torch
from compressai.zoo import models
from PIL import Image
import numpy as np

MAGIC = b"WAI1"


def log2_exact(n: int) -> int:
    e = int(round(math.log2(n)))
    if (1 << e) != n:
        raise ValueError(f"latent side {n} is not a power of two; decoder uses S_log2")
    return e


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--quality", type=int, default=3)
    ap.add_argument("--size", type=int, default=256, help="square crop/resize side")
    ap.add_argument("--image", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--model-dir", required=True,
                    help="dir with the EXPORTED hyper_synthesis.onnx — the encoder MUST compute "
                         "the entropy params with the same h_s the decoder runs, or float "
                         "round-off flips scale buckets and desyncs the rANS stream.")
    args = ap.parse_args()

    net = models["mbt2018-mean"](quality=args.quality, pretrained=True).eval()
    net.update(force=True)
    gc = net.gaussian_conditional

    img = Image.open(args.image).convert("RGB").resize((args.size, args.size), Image.LANCZOS)
    x = torch.from_numpy(np.asarray(img, dtype=np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0)

    import onnxruntime as ort
    hs = ort.InferenceSession(str(Path(args.model_dir) / "hyper_synthesis.onnx"),
                              providers=["CPUExecutionProvider"])

    with torch.no_grad():
        y = net.g_a(x)
        z = net.h_a(y)
        z_strings = net.entropy_bottleneck.compress(z)
        z_hat = net.entropy_bottleneck.decompress(z_strings, z.size()[-2:])
        # Entropy params from the EXPORTED ONNX h_s — identical to the decoder.
        in_name = hs.get_inputs()[0].name
        params = torch.from_numpy(hs.run(None, {in_name: z_hat.numpy()})[0])
        scales, means = params.chunk(2, 1)
        indexes = gc.build_indexes(scales)
        y_strings = gc.compress(y, indexes, means=means)
    y_bytes = y_strings[0]
    z_bytes = z_strings[0]
    _, Cy, Sy, _ = y.shape
    _, Cz, Sz, _ = z.shape
    if len(z_bytes) % 4 or len(y_bytes) % 4:
        raise RuntimeError("rANS bitstream not a multiple of 4 bytes — decoder seeds 32-bit words")

    H = W = args.size
    header = struct.pack("<IIIIHHBBH", H, W, len(z_bytes), len(y_bytes),
                         Cz, Cy, log2_exact(Sz), log2_exact(Sy), args.quality)
    payload = header + z_bytes + y_bytes

    manifest = {
        "wai": "1.0",
        "media": "image",
        "intent": "replicate",
        "model_requirement": {"capability": "wai.neural.mbt2018_mean"},
        "conditioning": {"kind": "mbt2018_mean"},
        "target": {"w": W, "h": H},
    }
    mb = json.dumps(manifest, separators=(",", ":")).encode()
    blob = MAGIC + struct.pack("<I", len(mb)) + mb + struct.pack("<I", len(payload)) + payload
    Path(args.out).write_bytes(blob)
    print(f"wrote {len(blob):,} B → {args.out}  "
          f"(z={len(z_bytes)}B, y={len(y_bytes)}B, latents y[1,{Cy},{Sy},{Sy}] z[1,{Cz},{Sz},{Sz}])")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())