"""Export sink artifacts for the `wai.neural.mlicpp` channel-autoregressive
codec decoded by `wai-rs/src/neural/mlicpp.rs`.

mlicpp.rs needs (next to the registered model path):
  hyper_synthesis.onnx   z_hat                  -> [1, 2*Cy, Sy, Sy]  (base scales ‖ means)
  entropy.onnx           concat[y_sofar, hyper] -> [1, 2*Cy, Sy, Sy]  (refined params)
                         input is [1, 3*Cy, Sy, Sy] (y is Cy ch, hyper is 2*Cy)
  synthesis.onnx         y_hat                  -> [1, 3, H, W]
  hyper_cdfs_q<q>.json   factorized z CDF (EntropyBottleneck)
  gauss_cdfs_q<q>.json   GaussianConditional CDF table

It wraps a checkpoint exposing `h_s`, `entropy` (the multi-reference net in the
WAI-native single-net form: concat[y_sofar, hyper] → params), `g_s`,
`entropy_bottleneck`, `gaussian_conditional`. This is the MLIC++ *class*
(channel-autoregressive multi-reference); the attention internals live inside
`entropy.onnx`. CDF extraction matches the hyperprior tool (same coders).

Run (needs torch + onnxruntime; bring your own mlicpp-class checkpoint):
  python tools/wai_mlicpp_export.py --ckpt mlic.pt --quality 3 --size 256 \
    --out-dir wai-web/demo/models/mlicpp
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch

from wai_mbt2018_mean_export import factorized_cdf_blob, gaussian_cdf_blob  # noqa: E402


def export_onnx(module, dummy, path: Path, opset):
    torch.onnx.export(
        module, (dummy,), path.as_posix(),
        input_names=["i"], output_names=["o"],
        dynamic_axes={"i": {2: "h", 3: "w"}, "o": {2: "h", 3: "w"}},
        opset_version=opset, do_constant_folding=True, dynamo=False)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True, help="mlicpp-class ckpt: h_s/entropy/g_s + entropy coders")
    ap.add_argument("--quality", type=int, default=3)
    ap.add_argument("--size", type=int, default=256)
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--opset", type=int, default=17)
    args = ap.parse_args()
    out = Path(args.out_dir); out.mkdir(parents=True, exist_ok=True)
    q = args.quality

    net = torch.load(args.ckpt, map_location="cpu", weights_only=False).eval()
    net.update(force=True)

    with torch.no_grad():
        y = net.g_a(torch.zeros(1, 3, args.size, args.size))
        z = net.h_a(y)
    _, M, Sy, _ = y.shape
    _, N, Sz, _ = z.shape
    print(f"  y[1,{M},{Sy},{Sy}]  z[1,{N},{Sz},{Sz}]  entropy-in=[1,{3 * M},{Sy},{Sy}]")

    export_onnx(net.h_s, torch.zeros(1, N, Sz, Sz), out / "hyper_synthesis.onnx", args.opset)
    export_onnx(net.entropy, torch.zeros(1, 3 * M, Sy, Sy), out / "entropy.onnx", args.opset)
    export_onnx(net.g_s, torch.zeros(1, M, Sy, Sy), out / "synthesis.onnx", args.opset)

    (out / f"hyper_cdfs_q{q}.json").write_text(
        json.dumps(factorized_cdf_blob(net.entropy_bottleneck, q), separators=(",", ":")))
    (out / f"gauss_cdfs_q{q}.json").write_text(
        json.dumps(gaussian_cdf_blob(net.gaussian_conditional, q), separators=(",", ":")))
    print(f"wrote 3 ONNX + 2 CDF tables → {out}")
    return 0


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