"""Encode an image to a `wai.neural.elic` envelope, decoded by
`wai-rs/src/neural/context.rs`.

The decoder's symbol order is the contract: ONE y rANS stream, read as
  [ slice over ALL channels, raster, where (y+x) even ]  (anchors)
  then
  [ same, where (y+x) odd ]                              (non-anchors)
channel-major. We flatten the symbols + scale-bucket indexes in exactly that
order and hand them to CompressAI's rANS in one shot, so the bytes the Rust
RansDecoder reads line up symbol-for-symbol. Anchors use the hyperprior base
params; non-anchors use the context-net params (the net sees the quantized
anchors). z uses the factorized EntropyBottleneck.

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

Run (needs torch + compressai-style ELIC ckpt + Pillow):
  python tools/wai_elic_encode.py --ckpt elic.pt --quality 3 --size 256 \
    --image in.png --out out.elic.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

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} 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("--image", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--model-dir", required=True,
                    help="dir with the EXPORTED hyper_synthesis.onnx + context.onnx — the encoder "
                         "MUST compute entropy params with the SAME nets the decoder runs (a "
                         "torch-vs-ONNX float drift flips scale buckets and desyncs the rANS; "
                         "verified on mbt2018_mean).")
    args = ap.parse_args()

    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
    from pathlib import Path as _P
    hs = ort.InferenceSession(str(_P(args.model_dir) / "hyper_synthesis.onnx"), providers=["CPUExecutionProvider"])
    ctxnet = ort.InferenceSession(str(_P(args.model_dir) / "context.onnx"), providers=["CPUExecutionProvider"])
    onnx = lambda s, t: torch.from_numpy(s.run(None, {s.get_inputs()[0].name: t.numpy()})[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:])
        base = onnx(hs, z_hat)                              # ONNX h_s — matches the decoder
        scales_b, means_b = base.chunk(2, 1)

        # Anchors first (need them to drive the context net).
        _, Cy, Sy, _ = y.shape
        yy, xx = torch.meshgrid(torch.arange(Sy), torch.arange(Sy), indexing="ij")
        anchor = ((yy + xx) % 2 == 0)                      # [Sy, Sy] bool
        sym_b = torch.round(y - means_b).to(torch.int32)   # anchor symbols
        y_anchor = ((sym_b.float() + means_b) * anchor).to(y.dtype)  # quantized anchors, rest 0
        ctx = onnx(ctxnet, y_anchor)                       # ONNX context — matches the decoder
        scales_c, means_c = ctx.chunk(2, 1)
        sym_c = torch.round(y - means_c).to(torch.int32)   # non-anchor symbols

    idx_b = gc.build_indexes(scales_b)[0]   # [Cy, Sy, Sy]
    idx_c = gc.build_indexes(scales_c)[0]
    sym_b, sym_c = sym_b[0], sym_c[0]
    anc = anchor.numpy()

    # Flatten in the decoder's order: anchors (chan-major raster), then non.
    syms, idxs = [], []
    for is_anchor in (True, False):
        s_t = sym_b if is_anchor else sym_c
        i_t = idx_b if is_anchor else idx_c
        for ch in range(Cy):
            for r in range(Sy):
                for c in range(Sy):
                    if bool(anc[r, c]) == is_anchor:
                        syms.append(int(s_t[ch, r, c]))
                        idxs.append(int(i_t[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("<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.elic"},
                "conditioning": {"kind": "elic"}, "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}  (z={len(z_bytes)}B y={len(y_bytes)}B, "
          f"{len(syms)} symbols)")
    return 0


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