"""Export the mean-scale HYPERPRIOR (mbt2018-mean) sink artifacts for the
`wai.neural.mbt2018_mean` capability decoded by `wai-rs/src/neural/hyperprior.rs`.

Unlike bmshj2018 (factorized, one static CDF per channel), the hyperprior
conditions each main latent `y` on a side stream `z`: the sink decodes `z`,
runs `h_s(z)` to get per-element (scale, mean), and decodes `y` with a
Gaussian-conditional model. So the sink needs FOUR files, all written next to
the registered model path:

  hyper_synthesis.onnx   z_hat  -> [1, 2*M, Sy, Sy]   (scales ‖ means, CompressAI chunk order)
  synthesis.onnx         y_hat  -> [1, 3,  H,  W]      (in [0,1])
  hyper_cdfs_q<q>.json    factorized z CDF table       (EntropyBottleneck)
  gauss_cdfs_q<q>.json    { scale_table, cdf[], cdf_length[], offset[] }  (GaussianConditional)

The encoder side (wai_mbt2018_mean_encode.py) produces the WAI envelope whose
two rANS bitstreams these tables decode. CompressAI's rANS is the same coder
`wai-rs/src/neural/rans.rs` ports (precision=16, bypass=4), so the bytes match.

Run (needs torch + compressai + onnxruntime):
  python tools/wai_mbt2018_mean_export.py --quality 3 --size 256 \
    --out-dir wai-web/demo/models/mbt2018_mean
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import torch
from compressai.zoo import models


class HyperSynthesis(torch.nn.Module):
    """z_hat -> gaussian params [1, 2*M, Sy, Sy] (scales ‖ means)."""
    def __init__(self, net):
        super().__init__()
        self.h_s = net.h_s

    def forward(self, z_hat: torch.Tensor) -> torch.Tensor:
        return self.h_s(z_hat)


class Synthesis(torch.nn.Module):
    """y_hat -> image [1, 3, H, W] in [0,1]."""
    def __init__(self, net):
        super().__init__()
        self.g_s = net.g_s

    def forward(self, y_hat: torch.Tensor) -> torch.Tensor:
        return self.g_s(y_hat).clamp(0.0, 1.0)


def factorized_cdf_blob(entropy_bottleneck, quality: int) -> dict:
    cdf = entropy_bottleneck._quantized_cdf.cpu().int().tolist()
    cdf_len = entropy_bottleneck._cdf_length.cpu().int().tolist()
    offset = entropy_bottleneck._offset.cpu().int().tolist()
    medians = entropy_bottleneck._get_medians().detach().reshape(-1).cpu().float().tolist()
    return {
        "model": "mbt2018-mean/hyper-z",
        "quality": quality,
        "precision": 16, "bypass_precision": 4,
        "channels": len(cdf_len),
        "cdf_length": cdf_len,
        "offset": offset,
        "medians": medians,
        "cdf": [row[:cdf_len[i]] for i, row in enumerate(cdf)],
    }


def gaussian_cdf_blob(gc, quality: int) -> dict:
    cdf = gc._quantized_cdf.cpu().int().tolist()
    cdf_len = gc._cdf_length.cpu().int().tolist()
    offset = gc._offset.cpu().int().tolist()
    scale_table = gc.scale_table.detach().reshape(-1).cpu().float().tolist()
    return {
        "model": "mbt2018-mean/gaussian-conditional",
        "quality": quality,
        "precision": 16, "bypass_precision": 4,
        "scale_table": scale_table,
        "cdf_length": cdf_len,
        "offset": offset,
        "cdf": [row[:cdf_len[i]] for i, row in enumerate(cdf)],
    }


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--quality", type=int, default=3, help="CompressAI quality 1..8; must match encoder")
    ap.add_argument("--size", type=int, default=256, help="square image side; probes latent shapes")
    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

    print(f"loading mbt2018-mean (quality={q})…")
    net = models["mbt2018-mean"](quality=q, pretrained=True).eval()
    net.update(force=True)

    # Probe latent shapes at the chosen image size.
    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"  main y: [1,{M},{Sy},{Sy}]   hyper z: [1,{N},{Sz},{Sz}]")

    z_dummy = torch.zeros(1, N, Sz, Sz)
    y_dummy = torch.zeros(1, M, Sy, Sy)
    hs, gs = HyperSynthesis(net).eval(), Synthesis(net).eval()
    with torch.no_grad():
        params = hs(z_dummy)
        assert params.shape[1] == 2 * M, f"h_s gave {params.shape[1]} ch, want 2*M={2*M}"
        img = gs(y_dummy)
    print(f"  h_s -> {tuple(params.shape)} (scales‖means)   g_s -> {tuple(img.shape)}")

    print("exporting hyper_synthesis.onnx + synthesis.onnx…")
    torch.onnx.export(
        hs, (z_dummy,), (out / "hyper_synthesis.onnx").as_posix(),
        input_names=["z"], output_names=["params"],
        dynamic_axes={"z": {2: "Sz", 3: "Sz"}, "params": {2: "Sy", 3: "Sy"}},
        opset_version=args.opset, do_constant_folding=True, dynamo=False,
    )
    torch.onnx.export(
        gs, (y_dummy,), (out / "synthesis.onnx").as_posix(),
        input_names=["y"], output_names=["image"],
        dynamic_axes={"y": {2: "Sy", 3: "Sy"}, "image": {2: "H", 3: "W"}},
        opset_version=args.opset, do_constant_folding=True, dynamo=False,
    )

    (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 2 ONNX + hyper_cdfs_q{q}.json + gauss_cdfs_q{q}.json → {out}")

    # ONNX-vs-torch sanity.
    try:
        import numpy as np
        import onnxruntime as ort
        s_hs = ort.InferenceSession((out / "hyper_synthesis.onnx").as_posix(),
                                    providers=["CPUExecutionProvider"])
        d = float(np.max(np.abs(s_hs.run(None, {"z": z_dummy.numpy()})[0] - params.numpy())))
        print(f"onnxruntime h_s max|onnx-torch| = {d:.3e}")
        if d > 1e-3:
            print("WARNING: h_s numerical mismatch > 1e-3", file=sys.stderr)
    except Exception as e:  # noqa: BLE001
        print(f"(skipped onnxruntime check: {e})")
    return 0


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