#!/usr/bin/env python3
"""Export the pretrained mbt2018_mean mean-scale hyperprior (CompressAI zoo,
quality 6, MSE) as an INTEGER-exact decoder and measure what that costs:
per-channel int12 weights + fixed-point activations + integer IGDN, all i64 —
byte-identical on every machine (NeuralIntegerExact). The tool reports the
integer and float PSNR on the given image; tests/int_hyper_synth.rs requires the
integer decode within 1 dB of the float decode.

  python tools/wai_sota_int_synth_export.py <image.png> <out-dir>

Normally run through tools/wai_sota_vectors.sh, which pins the source image and
the Python stack and produces the whole vector directory the Rust int_hyper tests
read.
"""
import sys, os, io, json, hashlib
import numpy as np, torch, warnings; warnings.filterwarnings('ignore')
from compressai.zoo import image as zoo_image
from compressai.zoo.pretrained import load_pretrained

AFRAC = 12; WBITS = 12; SHIFT = 24; GF = 16

# The pretrained checkpoint the model-derived vectors come from (all but the source
# image, source.png and source.rgb), pinned by its full SHA-256.
# CompressAI's own `mbt2018_mean(pretrained=True)` fetches it through torch.hub
# without checking any hash (and unpickles it with weights_only=False), so a changed
# object at the URL, or a stale file in the torch.hub cache, would silently yield
# different vectors that the Rust tests — comparing against Python output built
# from the same checkpoint — would still pass. So the bytes are checked here, and
# only then loaded, as tensors only (weights_only=True).
CKPT_URL = "https://compressai.s3.amazonaws.com/models/v1/mbt2018-mean-6-a19628ab.pth.tar"
CKPT_SHA256 = "a19628aba19732125e02caf5c38e233edbe2a641145a5b1779c51f4527f51dd6"


def load_pinned_mbt2018_mean():
    """mbt2018_mean quality 6 (MSE), built from checkpoint bytes whose SHA-256 is
    CKPT_SHA256 — or exit non-zero. Downloads into the torch.hub cache if absent."""
    zoo_url = zoo_image.model_urls["mbt2018-mean"]["mse"][6]
    if zoo_url != CKPT_URL:
        sys.exit(f"error: the installed compressai points mbt2018_mean q6 at {zoo_url}, "
                 f"not the pinned {CKPT_URL}; review the new checkpoint before re-pinning")
    path = os.path.join(torch.hub.get_dir(), "checkpoints", os.path.basename(CKPT_URL))
    if not os.path.exists(path):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        torch.hub.download_url_to_file(CKPT_URL, path, progress=False)
    with open(path, "rb") as f:
        blob = f.read()
    got = hashlib.sha256(blob).hexdigest()
    if got != CKPT_SHA256:
        sys.exit(f"error: {path} has sha256 {got}, pinned {CKPT_SHA256}. "
                 "Delete the file to re-download; if the published checkpoint itself "
                 "changed, re-pin only after reviewing it.")
    print(f"checkpoint {os.path.basename(path)}: sha256 {got} (pinned, verified)")
    state = torch.load(io.BytesIO(blob), map_location="cpu", weights_only=True)
    return zoo_image.model_architectures["mbt2018-mean"].from_state_dict(load_pretrained(state))

def q_weight_pc(w):  # ConvTranspose [in,out,k,k] → per-OUT-channel int12 + scale
    s = np.abs(w).max(axis=(0,2,3), keepdims=True)/(2**(WBITS-1)-1) + 1e-12
    wi = np.round(w/s).clip(-(2**(WBITS-1)), 2**(WBITS-1)-1).astype(np.int64)
    return wi, s.reshape(-1)

def q_weight_conv(w):  # Conv2d [out,in,k,k] → per-OUT-channel int12 + scale
    s = np.abs(w).max(axis=(1,2,3), keepdims=True)/(2**(WBITS-1)-1) + 1e-12
    wi = np.round(w/s).clip(-(2**(WBITS-1)), 2**(WBITS-1)-1).astype(np.int64)
    return wi, s.reshape(-1)

def conv_int(xi, wi, mult, bias, stride, pad, shift):
    """Integer regular conv (stride/pad) with per-out-channel (mult,shift) requant —
    the EXACT arithmetic of Rust `int_hyper_synth` 'conv'. `wi` is [out,in,k,k]."""
    co, ci, k, _ = wi.shape; H, W = xi.shape[1:]
    Ho = (H + 2*pad - k)//stride + 1; Wo = (W + 2*pad - k)//stride + 1
    xp = np.zeros((ci, H+2*pad, W+2*pad), dtype=np.int64); xp[:, pad:pad+H, pad:pad+W] = xi
    raw = np.zeros((co, Ho, Wo), dtype=np.int64)
    for ky in range(k):
        for kx in range(k):
            sub = xp[:, ky:ky+stride*Ho:stride, kx:kx+stride*Wo:stride]   # [ci,Ho,Wo]
            raw += np.tensordot(wi[:, :, ky, kx], sub, axes=([1], [0]))   # [co,Ho,Wo]
    return ((raw.astype(object)*mult[:,None,None]) >> shift).astype(np.int64) + bias[:,None,None]

def conv_t_int(x_int, wi, scale_oc, bias, stride, pad, opad):
    cin, co, k, _ = wi.shape; H, W = x_int.shape[1:]
    Ho = (H-1)*stride + k + opad - 2*pad; Wo = (W-1)*stride + k + opad - 2*pad
    raw = np.zeros((co, Ho+2*pad, Wo+2*pad), dtype=np.int64)
    for ic in range(cin):
        xi = x_int[ic]
        for (iy, ix) in np.argwhere(xi != 0):
            raw[:, iy*stride:iy*stride+k, ix*stride:ix*stride+k] += wi[ic] * xi[iy, ix]
    raw = raw[:, pad:pad+Ho, pad:pad+Wo]
    mult = np.round(scale_oc*(2**SHIFT)).astype(np.int64)
    bias_int = np.round(bias*(2**AFRAC)).astype(np.int64)
    return ((raw*mult[:,None,None]) >> SHIFT) + bias_int[:,None,None]

def isqrt_vec(n):  # exact floor(sqrt(n)) for int64 ≥0 — matches Rust u128::isqrt
    s = np.maximum(np.floor(np.sqrt(n.astype(np.float64))).astype(np.int64), 0)
    s = np.where((s+1)*(s+1) <= n, s+1, s)
    s = np.where(s*s > n, s-1, s)
    return s

def igdn_int(x_int, gi, bi):   # gi [C,C] int@2^GF, bi [C] int@2^(2·AFRAC)
    C = x_int.shape[0]; xf = x_int.reshape(C,-1).astype(np.int64); x2 = xf*xf
    out = np.zeros_like(xf)
    for i in range(C):
        norm = np.maximum(bi[i] + ((gi[i][:,None]*x2) >> GF).sum(0), 0)
        out[i] = (xf[i]*isqrt_vec(norm)) >> AFRAC
    return out.reshape(x_int.shape)

def conv_t_raw(xi, wi, mult, bias, stride, pad, opad, shift):
    """Integer transposed conv with per-out-channel (mult,shift) requant — the
    EXACT arithmetic of Rust `int_hyper_synth::conv_t` (uses the exported int mult,
    not a recomputed float scale), so the Python expected == the Rust decode."""
    cin, co, k, _ = wi.shape; H, W = xi.shape[1:]
    Ho = (H-1)*stride + k + opad - 2*pad; Wo = (W-1)*stride + k + opad - 2*pad
    raw = np.zeros((co, Ho+2*pad, Wo+2*pad), dtype=np.int64)
    for ic in range(cin):
        xc = xi[ic]
        for (iy, ix) in np.argwhere(xc != 0):
            raw[:, iy*stride:iy*stride+k, ix*stride:ix*stride+k] += wi[ic] * xc[iy, ix]
    raw = raw[:, pad:pad+Ho, pad:pad+Wo]
    return ((raw.astype(object)*mult[:,None,None]) >> shift).astype(np.int64) + bias[:,None,None]

def run_int_net(model, xi):
    """Run an exported integer model dict (the same JSON Rust `run_net` consumes)
    on `xi` ([c,h,w] int64 @ 2^afrac). Mirrors Rust byte-for-byte."""
    shift = model['shift']
    for L in model['layers']:
        t = L['type']
        if t == 'convT':
            cin, co, k = L['in_ch'], L['out_ch'], L['k']
            wi = np.array(L['weight'], dtype=np.int64).reshape(cin, co, k, k)
            mult = np.array(L['mult'], dtype=np.int64); bias = np.array(L['bias'], dtype=np.int64)
            xi = conv_t_raw(xi, wi, mult, bias, L['stride'], L['pad'], L['opad'], shift)
        elif t == 'conv':
            cin, co, k = L['in_ch'], L['out_ch'], L['k']
            wi = np.array(L['weight'], dtype=np.int64).reshape(co, cin, k, k)
            mult = np.array(L['mult'], dtype=np.int64); bias = np.array(L['bias'], dtype=np.int64)
            xi = conv_int(xi, wi, mult, bias, L['stride'], L['pad'], shift)
        elif t == 'igdn':
            C = L['c']; gi = np.array(L['gamma'], dtype=np.int64).reshape(C, C); bi = np.array(L['beta'], dtype=np.int64)
            xi = igdn_int(xi, gi, bi)
        elif t == 'relu':
            xi = np.maximum(xi, 0)
        elif t == 'leaky':
            xi = np.where(xi >= 0, xi, (xi*L['slope_q']) >> L['slope_frac'])
    return xi

def det_build_index_vec(sigma_q, edges):
    """Integer σ→bucket, vectorized — the EXACT logic of Rust
    `int_entropy::det_build_index` (and CompressAI `build_indexes`, but on the
    integer σ so encoder and Rust decoder agree by construction)."""
    n = len(edges); sq = np.maximum(sigma_q, edges[0])
    idx = np.full(sq.shape, n-1, dtype=np.int64)
    for e in edges[:n-1]:
        idx -= (sq <= e).astype(np.int64)
    return idx

def fnv1a(buf):
    h = 14695981039346656037
    for byte in buf:
        h = ((h ^ byte) * 1099511628211) & 0xFFFFFFFFFFFFFFFF
    return h

def model_to_bin(model):
    """Pack an integer model dict into the compact binary `IntModel::from_bytes`
    reads: magic WHM1, afrac/gf/shift/n_layers (u32), then per layer a u8 tag with
    int12 weights as i16 and IGDN gamma as i32 (the rest i64) — roughly halves the
    on-disk/wire size vs JSON decimal, so the in-browser model download is practical."""
    import struct
    out = bytearray(b'WHM1')
    out += struct.pack('<IIII', model['afrac'], model['gf'], model['shift'], len(model['layers']))
    for L in model['layers']:
        t = L['type']
        if t == 'convT':
            out += struct.pack('<BIIIIII', 0, L['in_ch'], L['out_ch'], L['k'], L['stride'], L['pad'], L['opad'])
            w = np.array(L['weight'], dtype=np.int64); assert np.abs(w).max() < 2**15, 'convT weight > i16'
            out += w.astype('<i2').tobytes()
            out += np.array(L['mult'], dtype='<i8').tobytes() + np.array(L['bias'], dtype='<i8').tobytes()
        elif t == 'conv':
            out += struct.pack('<BIIIII', 1, L['in_ch'], L['out_ch'], L['k'], L['stride'], L['pad'])
            w = np.array(L['weight'], dtype=np.int64); assert np.abs(w).max() < 2**15, 'conv weight > i16'
            out += w.astype('<i2').tobytes()
            out += np.array(L['mult'], dtype='<i8').tobytes() + np.array(L['bias'], dtype='<i8').tobytes()
        elif t == 'igdn':
            out += struct.pack('<BI', 2, L['c'])
            g = np.array(L['gamma'], dtype=np.int64); assert np.abs(g).max() < 2**31, 'gamma > i32'
            out += g.astype('<i4').tobytes() + np.array(L['beta'], dtype='<i8').tobytes()
        elif t == 'relu':
            out += struct.pack('<B', 3)
        elif t == 'leaky':
            out += struct.pack('<B', 4) + np.array([L['slope_q']], dtype='<i8').tobytes() + struct.pack('<I', L['slope_frac'])
    return bytes(out)

def main(imgpath, outdir):
    from PIL import Image
    # The ENCODER side (g_a, h_a, the float h_s used for the reference) is float, so
    # its latent can depend on kernel choice and reduction order. Both are pinned
    # to make two runs of this tool on one machine emit identical vectors. That
    # held for repeated runs on an arm64 macOS machine, but two runs on an x86_64
    # Linux runner wrote different cond.json and int_synth_sota.json, most likely
    # in psnr_float (`base` below), which the Rust tests read only as a quality
    # threshold; the latent and both rgb_hash values matched. The decode being
    # tested is integer and needs no such pinning.
    torch.use_deterministic_algorithms(True)
    torch.set_num_threads(1)
    # Verify the checkpoint before writing anything, so a refused checkpoint leaves
    # no partial vector directory behind.
    m = load_pinned_mbt2018_mean().eval()
    os.makedirs(outdir, exist_ok=True)
    img = Image.open(imgpath).convert('RGB')
    src8 = np.array(img)
    x = torch.from_numpy(src8.astype(np.float32)/255.).permute(2,0,1).unsqueeze(0)
    h,w = x.shape[2]//64*64, x.shape[3]//64*64; x = x[:,:,:h,:w]
    # The exact pixels the codec saw (HWC u8), so a decoder can measure its own
    # reconstruction quality instead of trusting the PSNR this tool reports.
    np.ascontiguousarray(src8[:h, :w]).tofile(f'{outdir}/source.rgb')
    def psnr(a,b): return float(-10*np.log10(((a-b)**2).mean()+1e-12))
    with torch.no_grad():
        y = m.g_a(x); z = m.h_a(y); z_hat = torch.round(z)
        gp = m.h_s(z_hat); _, means = gp.chunk(2,1)
        y_hat_t = torch.round(y - means) + means
        xr_float = m.g_s(y_hat_t).clamp(0,1)[0].numpy()
        y_hat = y_hat_t[0].numpy()
        y_np = y[0].numpy()  # raw float latent — the conditional encoder codes round(y - mu)
    base = psnr(x[0].numpy(), xr_float)

    xi = np.round(y_hat*(2**AFRAC)).astype(np.int64)
    model = {'afrac':AFRAC,'gf':GF,'shift':SHIFT,'layers':[]}
    inp = {'c':int(xi.shape[0]),'h':int(xi.shape[1]),'w':int(xi.shape[2]),'data':xi.reshape(-1).tolist()}
    for layer in m.g_s.children():
        t = type(layer).__name__
        if t == 'ConvTranspose2d':
            wi, sc = q_weight_pc(layer.weight.detach().numpy())
            mult = np.round(sc*(2**SHIFT)).astype(np.int64)
            bias_int = np.round(layer.bias.detach().numpy()*(2**AFRAC)).astype(np.int64)
            model['layers'].append({'type':'convT','in_ch':int(wi.shape[0]),'out_ch':int(wi.shape[1]),
                'k':int(wi.shape[2]),'stride':int(layer.stride[0]),'pad':int(layer.padding[0]),
                'opad':int(layer.output_padding[0]),'weight':wi.reshape(-1).tolist(),
                'mult':mult.tolist(),'bias':bias_int.tolist()})
            xi = conv_t_int(xi, wi, sc, layer.bias.detach().numpy(),
                            layer.stride[0], layer.padding[0], layer.output_padding[0])
        elif t == 'GDN':
            with torch.no_grad():
                beta = layer.beta_reparam(layer.beta).numpy(); C = beta.shape[0]
                gamma = layer.gamma_reparam(layer.gamma).reshape(C, C).numpy()
            gi = np.round(gamma*(2**GF)).astype(np.int64); bi = np.round(beta*(2**(2*AFRAC))).astype(np.int64)
            model['layers'].append({'type':'igdn','c':int(C),'gamma':gi.reshape(-1).tolist(),'beta':bi.tolist()})
            xi = igdn_int(xi, gi, bi)
    # INTEGER pixel quantization (deterministic, no float) — what the Rust port reproduces.
    out_c = np.clip(xi, 0, 1 << AFRAC)
    pix = (((out_c * 255 + (1 << (AFRAC-1))) >> AFRAC)).clip(0, 255).astype(np.uint8)  # [3,H,W] round-half-up
    rec = out_c.astype(np.float64) / (2**AFRAC)
    p_int = psnr(x[0].numpy(), rec); p_vf = psnr(xr_float, rec)
    print(f'FLOAT baseline {base:.2f} dB | INTEGER synthesis {p_int:.2f} dB (vs source) | {p_vf:.2f} dB (vs float synth)')
    rgb8 = np.ascontiguousarray(pix.transpose(1, 2, 0))  # HWC uint8
    # FNV-1a 64-bit over the pixel bytes — a dependency-free checksum the Rust
    # port reproduces with identical integer arithmetic (cross-language exact).
    h = 14695981039346656037
    for byte in rgb8.tobytes():
        h = ((h ^ byte) * 1099511628211) & 0xFFFFFFFFFFFFFFFF
    digest = format(h, '016x')
    json.dump({'afrac':AFRAC,'wbits':WBITS,'shift':SHIFT,'psnr_int':p_int,'psnr_float':base,'rgb_hash':digest,
               'out_c':int(rec.shape[0]),'out_h':int(rec.shape[1]),'out_w':int(rec.shape[2])},
              open(f'{outdir}/int_synth_sota.json','w'))
    json.dump(model, open(f'{outdir}/int_synth_sota_model.json','w'))
    open(f'{outdir}/int_synth_sota_model.bin','wb').write(model_to_bin(model))
    json.dump(inp, open(f'{outdir}/int_synth_sota_input.json','w'))
    print(f'wrote model ({len(model["layers"])} layers) + input + conformance (rgb_hash {digest[:16]}…) to {outdir}')

    # ---- integer hyper-decoder h_s: z_hat -> (sigma, mu), per-channel int12 + ReLU ----
    # z is entropy-coded with the FACTORIZED prior (CompressAI EntropyBottleneck) so the
    # whole bitstream is integer-entropy-coded, not part-raw. The channel median is folded
    # into a fixed-point constant median_q = round(median·2^AFRAC), so the decoder rebuilds
    # z_q = s·2^AFRAC + median_q with NO float — the integer-entropy contract on z too.
    from compressai.ans import RansEncoder as _RE
    eb = m.entropy_bottleneck; eb.update()
    with torch.no_grad():
        z = m.h_a(m.g_a(x))
    z_np = z[0].numpy()
    median = eb._get_medians().reshape(-1).detach().numpy().astype(np.float64)   # [Cz]
    median_q = np.round(median * (2**AFRAC)).astype(np.int64)
    z_sym = np.round(z_np - median[:,None,None]).astype(np.int64)                 # factorized symbols s
    z_cdfs = eb._quantized_cdf.tolist()
    z_cdf_lengths = eb._cdf_length.reshape(-1).int().tolist()
    z_offsets = eb._offset.reshape(-1).int().tolist()
    Cz = z_sym.shape[0]
    z_indexes = np.repeat(np.arange(Cz), z_sym.shape[1]*z_sym.shape[2])           # per-channel CDF
    z_bits = _RE().encode_with_indexes(z_sym.reshape(-1).tolist(), z_indexes.tolist(),
                                       z_cdfs, z_cdf_lengths, z_offsets)
    zi = z_sym * (2**AFRAC) + median_q[:,None,None]        # z_q: INTEGER latent @ 2^AFRAC (what the decoder rebuilds)
    hmodel = {'afrac':AFRAC,'gf':GF,'shift':SHIFT,'layers':[]}
    hin = {'c':int(zi.shape[0]),'h':int(zi.shape[1]),'w':int(zi.shape[2]),'data':zi.reshape(-1).tolist()}
    xi = zi.copy()
    for layer in m.h_s.children():
        t = type(layer).__name__
        if t == 'ConvTranspose2d':
            wi, sc = q_weight_pc(layer.weight.detach().numpy())
            mult = np.round(sc*(2**SHIFT)).astype(np.int64)
            bias_int = np.round(layer.bias.detach().numpy()*(2**AFRAC)).astype(np.int64)
            hmodel['layers'].append({'type':'convT','in_ch':int(wi.shape[0]),'out_ch':int(wi.shape[1]),
                'k':int(wi.shape[2]),'stride':int(layer.stride[0]),'pad':int(layer.padding[0]),
                'opad':int(layer.output_padding[0]),'weight':wi.reshape(-1).tolist(),
                'mult':mult.tolist(),'bias':bias_int.tolist()})
            xi = conv_t_int(xi, wi, sc, layer.bias.detach().numpy(), layer.stride[0], layer.padding[0], layer.output_padding[0])
        elif t == 'Conv2d':   # the final stride-1 conv → σ‖μ (M*3//2 → M*2)
            wi, sc = q_weight_conv(layer.weight.detach().numpy())
            mult = np.round(sc*(2**SHIFT)).astype(np.int64)
            bias_int = np.round(layer.bias.detach().numpy()*(2**AFRAC)).astype(np.int64)
            hmodel['layers'].append({'type':'conv','in_ch':int(wi.shape[1]),'out_ch':int(wi.shape[0]),
                'k':int(wi.shape[2]),'stride':int(layer.stride[0]),'pad':int(layer.padding[0]),
                'weight':wi.reshape(-1).tolist(),'mult':mult.tolist(),'bias':bias_int.tolist()})
            xi = conv_int(xi, wi, mult, bias_int, layer.stride[0], layer.padding[0], SHIFT)
        elif t == 'ReLU':
            hmodel['layers'].append({'type':'relu'})
            xi = np.maximum(xi, 0)
        elif t == 'LeakyReLU':  # slope 0.01 → fixed-point 655/2^16
            hmodel['layers'].append({'type':'leaky','slope_q':655,'slope_frac':16})
            xi = np.where(xi >= 0, xi, (xi*655) >> 16)
    sigmamu_int = xi.copy()            # [2M, sh, sw] integer σ‖μ @ 2^AFRAC, from the INTEGER h_s
    # FNV-1a over the integer sigma||mu output (LE i64 bytes) — matches Rust run_net
    hh = fnv1a(xi.astype('<i8').tobytes())
    json.dump(hmodel, open(f'{outdir}/int_hs_model.json','w'))
    open(f'{outdir}/int_hs_model.bin','wb').write(model_to_bin(hmodel))
    json.dump(hin, open(f'{outdir}/int_hs_input.json','w'))
    json.dump({'afrac':AFRAC,'out_c':int(xi.shape[0]),'out_h':int(xi.shape[1]),'out_w':int(xi.shape[2]),
               'params_hash':format(hh,'016x')}, open(f'{outdir}/int_hs.json','w'))
    print(f'wrote h_s ({len(hmodel["layers"])} layers): z_hat {zi.shape} -> sigma||mu {xi.shape}, params_hash {format(hh,"016x")}')

    # ---- the conditional ENTROPY layer: code y against integer σ (the bitstream) ----
    # Everything the DECODER needs is integer: z_hat (stored) → integer h_s → σ,μ_int;
    # σ_int → bucket (int_entropy); rANS-decode residual r; y_hat_int = r·2^AFRAC + μ_int.
    # The encoder runs the SAME integer h_s so it agrees on every σ-bucket — the whole
    # point of the integer-entropy contract. CDFs/offsets come from a real CompressAI
    # GaussianConditional; the rANS bytes from CompressAI's encoder (the Rust RansDecoder
    # is a byte-exact port of it, bypass path included), so no float touches the decode.
    from compressai.entropy_models import GaussianConditional
    from compressai.ans import RansEncoder
    M = sigmamu_int.shape[0] // 2
    sigma_int = sigmamu_int[:M]; mu_int = sigmamu_int[M:]          # each [M, sh, sw] @ 2^AFRAC
    scale_table = list(torch.exp(torch.linspace(np.log(0.11), np.log(256.0), 64)).numpy().astype(float))
    edges = np.floor(np.array(scale_table) * (2**AFRAC)).astype(np.int64)   # σ buckets @ 2^AFRAC
    gc = GaussianConditional(scale_table); gc.update()
    cdfs = gc._quantized_cdf.tolist()
    cdf_lengths = gc._cdf_length.reshape(-1).int().tolist()
    offsets = gc._offset.reshape(-1).int().tolist()
    # residual r = round(y - μ_int/2^AFRAC); decoder rebuilds y_hat_int = r·2^AFRAC + μ_int
    resid = np.round(y_np - mu_int.astype(np.float64) / (2**AFRAC)).astype(np.int64)
    indexes = det_build_index_vec(sigma_int, edges)                 # integer σ→bucket, per element
    enc = RansEncoder()
    bits = enc.encode_with_indexes(resid.reshape(-1).tolist(), indexes.reshape(-1).tolist(),
                                   cdfs, cdf_lengths, offsets)
    # expected reconstruction the DECODER must reproduce: integer g_s on y_hat_int
    y_hat_int = (resid << AFRAC) + mu_int
    rec_int = run_int_net(model, y_hat_int)
    out_c2 = np.clip(rec_int, 0, 1 << AFRAC)
    pix2 = (((out_c2 * 255 + (1 << (AFRAC-1))) >> AFRAC)).clip(0, 255).astype(np.uint8)
    rgb2 = np.ascontiguousarray(pix2.transpose(1, 2, 0))
    cond_hash = format(fnv1a(rgb2.tobytes()), '016x')
    p_cond = psnr(x[0].numpy(), out_c2.astype(np.float64) / (2**AFRAC))
    # entropy tables: the y Gaussian-conditional CDFs + the z FACTORIZED CDFs (+ median_q,
    # the fixed-point channel medians) — all integer, both bitstreams decode with no float.
    json.dump({'afrac':AFRAC,'frac':AFRAC,'scale_edges':edges.tolist(),
               'cdfs':cdfs,'cdf_lengths':cdf_lengths,'offsets':offsets,
               'z_cdfs':z_cdfs,'z_cdf_lengths':z_cdf_lengths,'z_offsets':z_offsets,
               'z_median_q':median_q.tolist()},
              open(f'{outdir}/cond_entropy.json','w'))
    json.dump({'m':int(M),'sh':int(sigma_int.shape[1]),'sw':int(sigma_int.shape[2]),
               'n_z':int(Cz),'zh':int(z_sym.shape[1]),'zw':int(z_sym.shape[2]),
               'z_rans':list(z_bits),'rans':list(bits)},
              open(f'{outdir}/cond_bitstream.json','w'))
    json.dump({'rgb_hash':cond_hash,'psnr_cond':p_cond,'psnr_float':base,
               'out_c':int(rec_int.shape[0]),'out_h':int(rec_int.shape[1]),'out_w':int(rec_int.shape[2])},
              open(f'{outdir}/cond.json','w'))
    print(f'wrote conditional bitstream: {len(z_bits)} z + {len(bits)} y rANS bytes '
          f'-> {p_cond:.2f} dB, rgb_hash {cond_hash}')


if __name__ == '__main__':
    main(sys.argv[1], sys.argv[2])