"""Encode an image to a `wai.neural.mlicpp` envelope, decoded by
`wai-rs/src/neural/mlicpp.rs`. This is the channel-autoregressive variant:
the latent is split into K channel slices decoded in order, each with a
checkerboard pass, and the entropy net is re-run on `concat[y_sofar, hyper]`
as the decoded context grows.

The decoder's symbol order is the contract (one y rANS stream):
  for slice k in 0..K:
    for pass in (anchors (y+x even), non-anchors (y+x odd)):
      params = entropy(concat[y_sofar, hyper]);  decode slice-k pass, chan-major
We replicate that order EXACTLY on the encode side — and crucially compute the
params with the EXPORTED ONNX `entropy`/`h_s` (not torch), so encoder and
decoder agree on scale buckets (a float drift desyncs the rANS; verified on
mbt2018_mean). z uses the factorized EntropyBottleneck.

Wire payload (== mlicpp.rs):
  <IIIIHHBBHH>  H, W, Lz, Ly, Cz, Cy, Sz_log2, Sy_log2, q, K  |  z bytes | y bytes

Run (needs torch + mlicpp-class ckpt + Pillow + onnxruntime):
  python tools/wai_mlicpp_encode.py --ckpt mlic.pt --quality 3 --size 256 \
    --slices 4 --image in.png --model-dir wai-web/demo/models/mlicpp \
    --out out.mlicpp.wai
"""
from __future__ import annotations

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

import numpy as np
import torch
from PIL import Image
from compressai.ans import BufferedRansEncoder


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


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--quality", type=int, default=3)
    ap.add_argument("--size", type=int, default=256)
    ap.add_argument("--slices", type=int, default=4, help="K channel slices (must divide Cy)")
    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 + entropy.onnx — encoder MUST "
                         "use the same nets the decoder runs or the rANS desyncs.")
    args = ap.parse_args()
    K = args.slices

    net = torch.load(args.ckpt, map_location="cpu", weights_only=False).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, 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"])
    en = ort.InferenceSession(str(Path(args.model_dir) / "entropy.onnx"), providers=["CPUExecutionProvider"])
    onnx = lambda s, t: torch.from_numpy(s.run(None, {s.get_inputs()[0].name: t.numpy().astype(np.float32)})[0])

    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:])
        hyper = onnx(hs, z_hat)
        _, Cy, Sy, _ = y.shape
        if Cy % K:
            raise SystemExit(f"Cy={Cy} not divisible by K={K}")
        sw = Cy // K
        yy, xx = torch.meshgrid(torch.arange(Sy), torch.arange(Sy), indexing="ij")
        anc = ((yy + xx) % 2 == 0).numpy()
        y_dec = torch.zeros(1, Cy, Sy, Sy)
        y0 = y[0]

        syms, idxs = [], []
        for k in range(K):
            c0, c1 = k * sw, (k + 1) * sw
            for is_anchor in (True, False):
                params = onnx(en, torch.cat([y_dec, hyper], 1))
                scales, means = params.chunk(2, 1)
                idx = gc.build_indexes(scales)[0]
                m0 = means[0]
                for ch in range(c0, c1):
                    for r in range(Sy):
                        for c in range(Sy):
                            if bool(anc[r, c]) == is_anchor:
                                sym = int(round(float(y0[ch, r, c] - m0[ch, r, c])))
                                syms.append(sym)
                                idxs.append(int(idx[ch, r, c]))
                                y_dec[0, ch, r, c] = sym + m0[ch, r, c]

    enc = BufferedRansEncoder()
    enc.encode_with_indexes(syms, idxs, gc._quantized_cdf.tolist(),
                            gc._cdf_length.tolist(), gc._offset.tolist())
    y_bytes = enc.flush()
    z_bytes = z_strings[0]
    if len(y_bytes) % 4 or len(z_bytes) % 4:
        raise RuntimeError("rANS bitstream not a multiple of 4 bytes")

    _, Cz, Sz, _ = z.shape
    H = W = args.size
    header = struct.pack("<IIIIHHBBHH", H, W, len(z_bytes), len(y_bytes),
                         Cz, Cy, log2_exact(Sz), log2_exact(Sy), args.quality, K)
    payload = header + z_bytes + y_bytes
    manifest = {"wai": "1.0", "media": "image", "intent": "replicate",
                "model_requirement": {"capability": "wai.neural.mlicpp"},
                "conditioning": {"kind": "mlicpp"}, "target": {"w": W, "h": H}}
    mb = json.dumps(manifest, separators=(",", ":")).encode()
    blob = b"WAI1" + 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}  (z={len(z_bytes)}B y={len(y_bytes)}B, "
          f"K={K} slices, {len(syms)} symbols)")
    return 0


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