#!/usr/bin/env python3
"""Extract JPEG AI model-CONSTANT tables into a model-asset bundle.

The self-sufficient decoder derives every *per-image* quantity from the
bitstream (the residual coding mask, the sigma surfaces, the indexes, the rvs
additions). What's left are *model constants* — fixed per checkpoint, identical
for every image — that a real decoder gets from the model it has installed:

  - the factorized entropy-bottleneck per-channel CDFs (the SOZ `pmfs`),
  - the per-channel dequant `scaler_vec` (the GainUnit divisor),
  - the dequant consts the route carried (z factorized-CDF offsets, lsbs tables).

The dequant gain is `resi * rvs_dq[grfs[ch]][lk(block)] / 2**16 / scaler_vec[ch]`
where `rvs_dq` (in model_constants) and `lk` (the per-block σ-skip "likely", the
same index the σ rvs uses) are derived in-decode; only the per-channel
`scaler_vec` is a model parameter (GainUnit's learned per-channel gain). This
tool extracts it once and verifies it reconstructs the whole gain map exactly.

These are architecturally image-invariant (learned model parameters), but the
reference's per-decode dump happened to emit them inside the per-image fixtures.
This tool lifts them out once into stable model-asset files, so the decode reads
them from the model bundle, not a per-image fixture.

  python3 tools/wai_jpegai_extract_model_consts.py <vm-assets-dir>

Writes into <vm-assets-dir>:
  jpegai_factorized_cdfs.json   {"y":[[...],...160], "uv":[[...],...96], "per_channel":16}
  jpegai_scaler_vec.json        {"y":[s0,...s159], "uv":[s0,...s95]}
  jpegai_dequant_consts.json    {"z_offsets":{"y":[...],"uv":[...]}, "lsbs_t0":[...], "lsbs_t1":[...]}
"""
import json, sys, os
import numpy as np

def main(base):
    ent = json.load(open(f"{base}/e2e/jpegai_real_entropy_fixture.json"))
    gains = json.load(open(f"{base}/e2e/dequant_gain_map_fixture.json"))
    route = json.load(open(f"{base}/e2e/route_fixture.json"))
    consts = json.load(open(f"{base}/model_constants.json"))
    adds = json.load(open(f"{base}/e2e/sigma_additions_fixture.json"))

    # --- factorized CDFs: the SOZ per-channel pmfs (model constant) ---
    soz = ent["SOZ"]["calls"]
    fac = {}
    for plane, call in zip(("y", "uv"), soz):
        ch = call["channels"]; plen = call["pmf_len"]
        flat = call["pmfs"]
        assert len(flat) == ch * plen, f"{plane}: {len(flat)} != {ch}*{plen}"
        fac[plane] = [flat[c*plen:(c+1)*plen] for c in range(ch)]
    fac["per_channel"] = soz[0]["per_channel"]
    json.dump(fac, open(f"{base}/jpegai_factorized_cdfs.json", "w"))

    # --- scaler_vec: the GainUnit per-channel divisor (model constant) ---
    # gain[ch][p] = rvs_dq[grfs[ch]][lk(block)] / 2**16 / scaler_vec[ch].
    # rvs_dq + lk are derived in-decode; only scaler_vec is a model parameter.
    # Extract it from one block and VERIFY it reconstructs every block exactly.
    dq = [np.array(consts["rvs_dq_t_flag0"]), np.array(consts["rvs_dq_t_flag1"])]
    dqn = len(dq[0])

    def block_lk(skip):  # per-block "likely" = ((Σ_8x8 skip + 32) >> 6), clamped
        lk = np.zeros(256, dtype=int)
        for y in range(16):
            for x in range(16):
                by, bx = (y // 8) * 8, (x // 8) * 8
                s = int(sum(int(skip[(by + yy) * 16 + bx + xx]) for yy in range(8) for xx in range(8)))
                lk[y * 16 + x] = min(max((s + 32) >> 6, 0), dqn - 1)
        return lk

    # grfs[ch] is the SAME per-channel flag the decode parses from the header and
    # uses for the σ rvs. Determine it unambiguously the way the σ rvs does — the
    # rvs_sigma table flag that reproduces the rvs addition — so scaler_vec is
    # extracted under the exact grfs the decode applies (not an empirical guess).
    sig0 = np.array(consts["rvs_sigma_t_flag0"]); sig1 = np.array(consts["rvs_sigma_t_flag1"])
    sv = {}
    for plane, base_call in (("y", 0), ("uv", 160)):
        g = np.array(gains[plane]["gain"]); km = np.array(gains[plane]["known_mask"])
        rvs_add = np.array(adds[plane]["rvs"]); nch = len(g) // 256
        vec = []
        for c in range(nch):
            fin = np.load(f"{base}/e2e/sigma/call{base_call + c}_scales.npy").reshape(-1)
            ra = rvs_add[c*256:(c+1)*256]
            lk = block_lk(fin - ra)
            # grfs = the rvs_sigma flag whose table reproduces this channel's rvs addition.
            grfs = 0 if np.array_equal(sig0[lk], ra) else 1
            assert np.array_equal((sig0 if grfs == 0 else sig1)[lk], ra), \
                f"{plane} ch{c}: neither rvs_sigma flag reproduces rvs_add"
            gc = g[c*256:(c+1)*256]; kp = np.where(km[c*256:(c+1)*256] > 0)[0]
            if len(kp) == 0:
                vec.append(1.0); continue
            s = dq[grfs][lk[kp[0]]] / 65536.0 / gc[kp[0]]
            assert np.abs(dq[grfs][lk[kp]] / 65536.0 / s - gc[kp]).max() < 1e-4, \
                f"{plane} ch{c}: scaler_vec did not reconstruct the gain map under parsed grfs"
            vec.append(float(s))
        sv[plane] = vec
    json.dump(sv, open(f"{base}/jpegai_scaler_vec.json", "w"))

    # --- dequant consts the route carried (all model constants) ---
    json.dump({"z_offsets": route["z_offsets"],
               "lsbs_t0": route["lsbs_t0"], "lsbs_t1": route["lsbs_t1"]},
              open(f"{base}/jpegai_dequant_consts.json", "w"))

    print(f"wrote jpegai_factorized_cdfs.json (y:{len(fac['y'])} uv:{len(fac['uv'])}), "
          f"jpegai_scaler_vec.json (y:{len(sv['y'])} uv:{len(sv['uv'])}), "
          f"jpegai_dequant_consts.json into {base}")

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