#!/usr/bin/env python3
"""Build every runtime input the whole-decoder test (`jpegai_e2e_full`) needs, from
the shipped JPEG-AI assets. Nothing here is derived or fitted — it is a straight
extraction of model weights and normative tables, so the end-to-end test is
reproducible rather than depending on hand-made scratch files.

    python wai_jpegai_build_e2e_inputs.py <assets_dir> <out_dir>

<assets_dir> is the jpegai-vm-assets root. Writes:
    <out_dir>/weights/{hs,hyp}_{y,uv}.json, mcm{0..3}_y.json, synth_{y,uv}.json, icci.json
    <out_dir>/norm.json   (pmf, bounds, per-channel z size, factorized CDFs, gain
                           vectors, LSBS tables, LEF op-point tables, reference RGB)

The weight exports call `wai_jpegai_export_weights.py` (same directory); norm.json is
assembled here. Requires numpy + Pillow (only for reading the reference PNG).

Model selection is signalled by the picture (icci_model_signalled_idx → Y=model_6,
UV=model_3); those two wrappers live under e2e/pfchain and are what the exporter's
`icci` mode reads.
"""
import json
import os
import subprocess
import sys

import numpy as np
from PIL import Image

HERE = os.path.dirname(os.path.abspath(__file__))
EXPORT = os.path.join(HERE, "wai_jpegai_export_weights.py")


def export(mode, *args):
    subprocess.run([sys.executable, EXPORT, mode, *args], check=True)


def main(assets, out):
    W = os.path.join(out, "weights")
    os.makedirs(W, exist_ok=True)
    M = os.path.join(assets, "model/CCS_SGMM/tools_0")
    M16 = os.path.join(assets, "model16/CCS_SGMM/tools_0")

    # --- model weights ---
    for pl in ("y", "uv"):
        export("int", f"{M}/model_{pl}/common_modules/hyper_scale_decoder.onnx", f"{W}/hs_{pl}.json")
        export("ptq", f"{M}/model_{pl}/common_modules/hyper_decoder.onnx", f"{W}/hyp_{pl}.json")
        export("ptq", f"{M16}/model_{pl}/synthesis16.onnx", f"{W}/synth_{pl}.json")
    for i in range(4):
        export("ptq", f"{M}/model_y/common_modules/MCM/stage{i}.onnx", f"{W}/mcm{i}_y.json")
    export("icci", f"{assets}/e2e/pfchain/icci_y_m6.onnx",
           f"{assets}/e2e/pfchain/icci_uv_m3.onnx", f"{W}/icci.json")

    # --- normative tables + per-image reference RGB ---
    fac = json.load(open(f"{assets}/jpegai_factorized_cdfs.json"))
    norm = json.load(open(f"{assets}/jpegai_normative_tables.json"))
    mc = json.load(open(f"{assets}/model_constants.json"))
    bounds = [int(x) for x in open(f"{assets}/me-tANS/bounds.csv") if x.strip()]

    def cdfs(pl):
        cd = fac[pl]
        return [int(x) for row in cd for x in row] if isinstance(cd[0], list) else [int(x) for x in cd]

    def gain(pl):
        p = f"{M}/model_{pl}/common_modules/quantizer/gain_unit_mlog.csv"
        return [int(v) for v in np.loadtxt(p, delimiter=",", dtype=int)]

    rgb = np.array(Image.open(f"{assets}/e2e/test256_dec.png").convert("RGB"))
    data = {
        "pmf": norm["pmf"], "bounds": bounds, "per_channel": fac["per_channel"],
        "cdfs_y": cdfs("y"), "cdfs_uv": cdfs("uv"), "gain_y": gain("y"), "gain_uv": gain("uv"),
        "lsbs_t0": [int(v) for v in np.load(f"{assets}/e2e/lsbs/lsbs_table0.npy")],
        "lsbs_t1": [int(v) for v in np.load(f"{assets}/e2e/lsbs/lsbs_table1.npy")],
        "lef_thr0": mc["lef_thr"][0], "lef_mag0": mc["lef_mag"][0],
        "ref_rgb": rgb.ravel().tolist(),
    }
    json.dump(data, open(f"{out}/norm.json", "w"))
    print(f"wrote {W}/*.json (11 weight files) and {out}/norm.json")


if __name__ == "__main__":
    if len(sys.argv) != 3:
        print(__doc__)
        sys.exit(1)
    main(sys.argv[1], sys.argv[2])