"""Golden-fixture harness for the Rust JPEG AI mANS port (jpegai_mans.rs).

Builds the ISO/IEC 6048-1 reference software mANS extension (BSD-3, public:
gitlab.com/wg1/jpeg-ai/jpeg-ai-reference-software) and produces
/tmp/jpegai_mans_goldens.json: decode tables built by the reference utils,
plus y (sgm, incl. escapes) and z (factorize, incl. escape-63) bitstreams
encoded by the REFERENCE encoder. The Rust test
`jpegai_mans_matches_reference_software` then byte-compares its own table
builder and decodes both streams. Build the ext first (clone repo to
/tmp/jpeg-ai-ref, omp shim, pybind11 — see tools/README or memory notes);
z values are 0..62 (63 is the internal escape marker).
"""
import sys, json
import numpy as np
sys.path.insert(0, '/tmp/jpegai_ext')
sys.path.insert(0, '/tmp/jpeg-ai-ref/src/codec/entropy_coding/lib_wrappers/mans')
import ans
from utils import get_cdf_matrix, get_encode_transitions, get_state_maps, get_decode_transitions

rng = np.random.RandomState(11)
# 4 distributions shaped like the normative set: short pdf rows (mass-256 pmfs
# over interleaved symbol order) + the mandatory 1 at index bound*2-1.
pdf_rows = [
    [128, 64, 32, 16, 8, 4, 2, 1],            # sharp, bound 4
    [64, 64, 48, 32, 24, 12, 8, 2, 1, 1],     # medium, bound 5
    [32, 32, 32, 32, 28, 24, 20, 16, 12, 8, 8, 4, 4, 2, 1, 1],  # wide, bound 8
    [16]*14 + [8]*2 + [4]*3 + [2]*1 + [1, 1],  # very wide, bound 11
]
bounds = np.array([4, 5, 8, 11], dtype=np.uint8)
NSYM = 256
pmf = np.zeros([4, NSYM], dtype=np.int64)
for i, row in enumerate(pdf_rows):
    row = list(row)
    s = sum(row)
    assert s <= 256, s
    row[0] += 256 - s          # normalize mass to exactly 256
    pmf[i, :len(row)] = row
    pmf[i, bounds[i]*2 - 1] = max(pmf[i, bounds[i]*2 - 1], 1)
    # re-fix mass after the bound-slot guarantee
    excess = pmf[i].sum() - 256
    pmf[i, 0] -= excess
cdf = get_cdf_matrix(pmf)
enc_t = get_encode_transitions(pmf, cdf)
dec_t = get_decode_transitions(pmf, cdf)
smaps = get_state_maps(cdf)

# --- y stream (sgm): 200 elements, some masked off, some outbound escapes ---
N = 200
sigma = rng.randint(0, 4, N).astype(np.uint8)
masks = (rng.rand(N) > 0.15)
r = np.zeros(N, dtype=np.int16)
for i in range(N):
    b = int(bounds[sigma[i]])
    if rng.rand() < 0.08:
        r[i] = rng.randint(b, 200) * (1 if rng.rand() < 0.5 else -1)  # escape
    else:
        r[i] = rng.randint(-(b-1), b)  # in-table
r[~masks] = 0
enc = ans.ANSEncoder(1 << 20, 1)
enc.set_sgm_transitions(enc_t.data, bounds.data, smaps.data)
r_enc = r.copy()
enc.encode_sgm(sigma, r_enc, masks)
size = enc.close()
mem = np.empty([size], dtype=np.uint8); enc.get_memory(mem)
ts = np.empty([1], dtype=np.uint32); enc.get_thread_sizes(ts)
off = np.cumsum(ts).astype(np.uint32)
# reference decode (sanity)
dec = ans.ANSDecoder(mem, off)
dec.set_sgm_transitions(dec_t.data, bounds.data)
r_dec = np.zeros(N, dtype=np.int16)
dec.decode_sgm(sigma, r_dec, masks)
assert (r_dec[masks] == r[masks]).all(), "reference round-trip failed"

# --- z stream (factorize): 2 channels x 96, incl. escape value 63 ---
def zcdf(rng):
    p = rng.randint(1, 9, 63).astype(np.int64)
    c = np.cumsum(p)
    return ((c * 255 + (c[-1] >> 1)) // c[-1]).astype(np.uint8)
zc = np.stack([zcdf(rng), zcdf(rng)])
z = rng.randint(0, 63, [2, 96]).astype(np.uint8)   # 0..62; 63 is the internal escape
enc2 = ans.ANSEncoder(1 << 20, 1)
enc2.set_sgm_transitions(enc_t.data, bounds.data, smaps.data)
z_enc = z.copy()
enc2.encode_factorize(zc, z_enc)
size2 = enc2.close()
mem2 = np.empty([size2], dtype=np.uint8); enc2.get_memory(mem2)
ts2 = np.empty([1], dtype=np.uint32); enc2.get_thread_sizes(ts2)
off2 = np.cumsum(ts2).astype(np.uint32)
dec2 = ans.ANSDecoder(mem2, off2)
dec2.set_sgm_transitions(dec_t.data, bounds.data)
z_dec = np.zeros_like(z)
dec2.decode_factorize(zc, z_dec)
assert (z_dec == z).all(), "z round-trip failed"

json.dump({
    "pmf": pmf.tolist(), "cdf": cdf.tolist(), "bounds": bounds.tolist(),
    "decode_transitions": dec_t.tolist(),
    "y": {"memory": mem.tolist(), "offsets": off.tolist(),
          "sigma": sigma.tolist(), "masks": masks.tolist(), "expected": r.tolist()},
    "z": {"memory": mem2.tolist(), "offsets": off2.tolist(),
          "cdfs": zc.tolist(), "expected": z.tolist()},
}, open("/tmp/jpegai_mans_goldens.json", "w"))
print(f"goldens OK: y stream {size}B ({int(masks.sum())} coded, escapes incl), z stream {size2}B (2ch x 96, escape 63 incl)")