//! Golden validation of the JPEG AI entropy decode against bitstreams produced by
//! the REFERENCE software's own encoder (ISO/IEC 6048-1 reference, BSD-3,
//! compiled from gitlab.com/wg1/jpeg-ai/jpeg-ai-reference-software), through the decode
//! the exact decoder and the float decode path both run: the integer operator set's
//! `mans_decode` (integer-operators §5E E8) and the scale graph (§13.6.1). The fixture
//! also carries the reference-built decode tables, so E8's table builder is
//! byte-compared too. Skipped if the fixture is absent (regenerate with the harness
//! documented in tools/wai_jpegai_goldens.py).
#![cfg(feature = "neural")]

use std::path::Path;

use wai_ops::stage::mans;

#[test]
fn jpegai_mans_matches_reference_software() {
    let fixture = Path::new("../jpegai-vm-assets/jpegai_mans_goldens.json");
    if !fixture.exists() {
        eprintln!("skipping jpegai mans golden: fixture not present");
        return;
    }
    let v: serde_json::Value =
        serde_json::from_slice(&std::fs::read(fixture).unwrap()).unwrap();
    let as_i64v = |x: &serde_json::Value| -> Vec<i64> {
        x.as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect()
    };

    // 1. E8's table builder reproduces the reference tables bit for bit.
    let pmf: Vec<Vec<i64>> = v["pmf"].as_array().unwrap().iter().map(&as_i64v).collect();
    let reference: Vec<Vec<i64>> =
        v["decode_transitions"].as_array().unwrap().iter().map(&as_i64v).collect();
    let built = mans::build_latent_tables(&pmf).expect("every row sums to 256");
    for (d, (b, r)) in built.iter().zip(&reference).enumerate() {
        for s in 0..256 {
            assert_eq!(b[s] as i64, r[s], "E8 table mismatch: distribution {d}, state {s}");
        }
    }
    println!("  wai-ops E8 table builder: {} distributions x 256 states bit-exact", built.len());

    // 2. The latent path (`sgm`) against the reference encoder's bitstream.
    let y = &v["y"];
    let mask: Vec<bool> = y["masks"].as_array().unwrap().iter().map(|m| m.as_bool().unwrap()).collect();
    let sigma: Vec<usize> = as_i64v(&y["sigma"]).iter().map(|&s| s as usize).collect();
    let mem: Vec<u8> = as_i64v(&y["memory"]).iter().map(|&b| b as u8).collect();
    let end = as_i64v(&y["offsets"])[0] as usize;
    let got = mans::decode_sgm(&mem[..end], &pmf, &as_i64v(&v["bounds"]), &[mans::Span { mask: mask.clone(), sigma }]).unwrap();
    let coded = mask.iter().filter(|&&m| m).count();
    for (i, &e) in as_i64v(&y["expected"]).iter().enumerate() {
        if mask[i] {
            assert_eq!(got[0][i], e as i16 as i64, "E8 y symbol {i} diverges");
        }
    }
    println!("  wai-ops E8 sgm: {coded} residuals (incl. escapes) match the reference encoder");

    // 3. The hyper-latent path (`fac`).
    let z = &v["z"];
    let mem: Vec<u8> = as_i64v(&z["memory"]).iter().map(|&b| b as u8).collect();
    let end = as_i64v(&z["offsets"])[0] as usize;
    let rows: Vec<Vec<i64>> = z["cdfs"].as_array().unwrap().iter().map(&as_i64v).collect();
    let want: Vec<Vec<i64>> = z["expected"].as_array().unwrap().iter().map(&as_i64v).collect();
    let lens: Vec<usize> = want.iter().map(Vec::len).collect();
    assert_eq!(mans::decode_fac(&mem[..end], &rows, &lens).unwrap(), want, "E8 z (fac) diverges from the reference encoder");
    println!("  wai-ops E8 fac: {} rows x {} hyper-latents match the reference encoder", want.len(), lens[0]);
}

/// End-to-end through all three layers: codestream framing → AE-substream unpacking →
/// E8's `fac` decode of the reference encoder's golden z-stream.
#[test]
fn jpegai_codestream_to_symbols() {
    use wai::neural::jpegai_bitstream::{parse_ae_substream, parse_codestream, MARKER_SOZ};

    let fixture = std::path::Path::new("../jpegai-vm-assets/jpegai_mans_goldens.json");
    if !fixture.exists() {
        eprintln!("skipping: fixture not present");
        return;
    }
    let v: serde_json::Value =
        serde_json::from_slice(&std::fs::read(fixture).unwrap()).unwrap();
    let as_i64v = |x: &serde_json::Value| -> Vec<i64> {
        x.as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect()
    };
    let z = &v["z"];
    let memory: Vec<u8> = as_i64v(&z["memory"]).iter().map(|&b| b as u8).collect();
    let rows: Vec<Vec<i64>> = z["cdfs"].as_array().unwrap().iter().map(&as_i64v).collect();
    let expected: Vec<Vec<i64>> = z["expected"].as_array().unwrap().iter().map(&as_i64v).collect();

    // Frame: SOC | SOZ + ue(size) + payload (single thread: no deltas) | EOC.
    let mut cs: Vec<u8> = 0xFF80u16.to_be_bytes().to_vec();
    cs.extend(MARKER_SOZ.to_be_bytes());
    let n = memory.len() as u64 + 1;
    let len = 64 - n.leading_zeros();
    let total_bits = (2 * len - 1) as usize;
    let padded = total_bits.div_ceil(8) * 8;
    let acc: u128 = (n as u128) << (padded - total_bits);
    for shift in (0..padded).step_by(8).rev() {
        cs.push(((acc >> shift) & 0xFF) as u8);
    }
    cs.extend_from_slice(&memory);
    cs.extend(0xFF81u16.to_be_bytes());

    let subs = parse_codestream(&cs).expect("codestream parses");
    assert_eq!(subs.len(), 1);
    assert_eq!(subs[0].marker, MARKER_SOZ);
    let (mem, offsets) = parse_ae_substream(subs[0].payload, 1).expect("ae unpacks");
    assert_eq!(offsets.len(), 1);
    let lens: Vec<usize> = expected.iter().map(Vec::len).collect();
    let got = mans::decode_fac(&mem[..offsets[0] as usize], &rows, &lens).unwrap();
    assert_eq!(got, expected, "end-to-end codestream → symbols diverged");
    println!("  codestream → AE unpack → wai-ops E8: {} symbols exact", lens.iter().sum::<usize>());
}

/// The NORMATIVE table set: 32 σ-distributions from the standard's literal
/// PMF constants (`pdf_r` in the reference `ec_lib_mans.py`), exported via
/// the reference software's own `store_parameters`. E8's table builder must
/// reproduce every packed word the runtime decoder consumes.
#[test]
fn jpegai_normative_tables_bit_exact() {
    let fixture = std::path::Path::new("../jpegai-vm-assets/jpegai_normative_tables.json");
    if !fixture.exists() {
        eprintln!("skipping: normative fixture not present");
        return;
    }
    let v: serde_json::Value =
        serde_json::from_slice(&std::fs::read(fixture).unwrap()).unwrap();
    let as_i64v = |x: &serde_json::Value| -> Vec<i64> {
        x.as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect()
    };
    let pmf: Vec<Vec<i64>> = v["pmf"].as_array().unwrap().iter().map(&as_i64v).collect();
    let packed: Vec<Vec<i64>> = v["packed"].as_array().unwrap().iter().map(&as_i64v).collect();
    assert_eq!(pmf.len(), 32);
    let built = mans::build_latent_tables(&pmf).expect("the normative rows sum to 256");
    for (d, (b, r)) in built.iter().zip(&packed).enumerate() {
        for s in 0..256 {
            assert_eq!(b[s] as i64, r[s], "E8 normative table mismatch: distribution {d}, state {s}");
        }
    }
    println!("  wai-ops E8: the normative tables, 32 distributions x 256 states bit-exact");
}

/// Parse a REAL bitstream produced by the reference encoder (256x256 RGB,
/// default operating point) — first contact with non-synthetic data.
#[test]
fn jpegai_real_bitstream_parses() {
    let path = std::path::Path::new("../jpegai-vm-assets/e2e/test256.bin");
    if !path.exists() {
        eprintln!("skipping: e2e fixture not present");
        return;
    }
    let buf = std::fs::read(path).unwrap();
    let subs = wai::neural::jpegai_bitstream::parse_codestream(&buf)
        .expect("real bitstream must parse");
    let names = |m: u16| match m {
        0xFF82 => "PIH", 0xFF83 => "TON", 0xFF84 => "RDI", 0xFF88 => "SOZ",
        0xFF89 => "SORP", 0xFF8A => "SORS", 0xFF8B => "SOQ", 0xFF8C => "UDI",
        _ => "?",
    };
    let layout: Vec<String> = subs.iter()
        .map(|s| format!("{}:{}", names(s.marker), s.payload.len()))
        .collect();
    println!("  real bitstream substreams: [{}]", layout.join(", "));
    assert!(subs.iter().any(|s| s.marker == 0xFF82), "must contain PIH");
    assert!(subs.iter().any(|s| matches!(s.marker, 0xFF88..=0xFF8B)), "must contain an AE substream");
    // every AE substream must unpack through the aemem layer
    for s in subs.iter().filter(|s| matches!(s.marker, 0xFF88..=0xFF8B)) {
        let unpacked = wai::neural::jpegai_bitstream::parse_ae_substream(s.payload, 1);
        assert!(unpacked.is_some(), "AE substream {} must unpack", names(s.marker));
    }
}

/// Replay every entropy-decode call the reference decoder made on the REAL
/// test256.bin bitstream through E8, with the normative tables: one `sgm` decode per
/// latent substream over all its channels' spans, and one `fac` decode over the
/// hyper-latent substream's rows, luma's then chroma's. Every one of the 258 calls'
/// symbols must be the reference's.
#[test]
fn jpegai_real_substreams_entropy_bit_exact() {
    let assets = std::path::Path::new("../jpegai-vm-assets");
    let fx_path = assets.join("e2e/jpegai_real_entropy_fixture.json");
    let norm_path = assets.join("jpegai_normative_tables.json");
    if !fx_path.exists() || !norm_path.exists() {
        eprintln!("skipping: real-entropy fixtures not present");
        return;
    }
    let norm: serde_json::Value =
        serde_json::from_slice(&std::fs::read(norm_path).unwrap()).unwrap();
    let pmf: Vec<Vec<i64>> = norm["pmf"].as_array().unwrap().iter()
        .map(|r| r.as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect())
        .collect();
    let tables = mans::build_latent_tables(&pmf).unwrap();
    let bounds: Vec<i64> = std::fs::read_to_string(assets.join("me-tANS/bounds.csv")).unwrap()
        .lines().filter(|l| !l.trim().is_empty())
        .map(|l| l.trim().parse().unwrap()).collect();

    let bin = std::fs::read(assets.join("e2e/test256.bin")).unwrap();
    let subs = wai::neural::jpegai_bitstream::parse_codestream(&bin).unwrap();
    let payload = |m: u16| subs.iter().find(|s| s.marker == m).unwrap().payload;
    let fx: serde_json::Value =
        serde_json::from_slice(&std::fs::read(&fx_path).unwrap()).unwrap();
    let mut calls = 0usize;
    for (name, marker) in [("SORP", 0xFF89u16), ("SORS", 0xFF8A), ("SOZ", 0xFF88)] {
        let (mem, ends) = wai::neural::jpegai_bitstream::parse_ae_substream(payload(marker), 1).unwrap();
        let st = &fx[name];
        let want_mem: Vec<u8> = st["mem"].as_array().unwrap().iter()
            .map(|n| n.as_u64().unwrap() as u8).collect();
        assert_eq!(mem, &want_mem[..], "{name}: unpacked memory must match dumped");
        let mem = &mem[..ends[0] as usize];
        let list = st["calls"].as_array().unwrap();
        let ints = |c: &serde_json::Value, k: &str| -> Vec<i64> { c[k].as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect() };
        if name == "SOZ" {
            let mut rows = vec![];
            let mut want = vec![];
            for call in list {
                let per = call["per_channel"].as_u64().unwrap() as usize;
                rows.extend(ints(call, "pmfs").chunks(63).map(<[i64]>::to_vec));
                want.extend(ints(call, "x").chunks(per).map(<[i64]>::to_vec));
            }
            let lens: Vec<usize> = want.iter().map(Vec::len).collect();
            assert_eq!(mans::decode_fac(mem, &rows, &lens).unwrap(), want, "{name}: E8 fac");
        } else {
            let spans: Vec<mans::Span> = list.iter().map(|c| mans::Span {
                mask: ints(c, "masks").iter().map(|&m| m != 0).collect(),
                sigma: ints(c, "indexes").iter().map(|&s| s as usize).collect(),
            }).collect();
            let got = mans::decode_sgm_tables(mem, &tables, &bounds, &spans).unwrap();
            for (i, (g, c)) in got.iter().zip(list).enumerate() {
                let want: Vec<i64> = ints(c, "x").iter().map(|&v| v as i16 as i64).collect();
                assert_eq!(g, &want, "{name}: E8 sgm span {i}");
            }
        }
        calls += list.len();
    }
    assert_eq!(calls, 258);
    println!("  real-bitstream entropy replay through wai-ops E8: {calls} calls bit-exact across SORP/SORS/SOZ");
}

/// The scale graph (integer-operators §13.6.1, lowered from the float path's bundle form
/// of the network): z (real, entropy-decoded) → the integer scale `I_s`, bit-exact
/// against the reference decoder's dumps for both planes. The graph is lowered with gains
/// of 0, so its output is `I_s` itself.
#[test]
fn jpegai_hyper_scale_decoder_int_bit_exact() {
    let assets = std::path::Path::new("../jpegai-vm-assets");
    let w_path = assets.join("hyper_scale_decoder_int_weights.json");
    let fx_path = assets.join("e2e/hyper_scale_io_fixture.json");
    if !w_path.exists() || !fx_path.exists() {
        eprintln!("skipping: hyper-scale fixtures not present");
        return;
    }
    let weights: serde_json::Value =
        serde_json::from_slice(&std::fs::read(&w_path).unwrap()).unwrap();
    let fx: serde_json::Value =
        serde_json::from_slice(&std::fs::read(&fx_path).unwrap()).unwrap();
    let ints = |x: &serde_json::Value| -> Vec<i64> {
        x.as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect()
    };
    for plane in ["y", "uv"] {
        let zs = &fx[plane];
        let shape: Vec<usize> = zs["z_shape"].as_array().unwrap().iter()
            .map(|n| n.as_u64().unwrap() as usize).collect();
        let (c, h, w) = (shape[1], shape[2], shape[3]);
        let want = ints(&zs["i_s"]);
        let net = wai::jpegai_ops::scale_net_of_blocks(&weights[plane]).unwrap();
        let g = wai::jpegai_ops::load_scale(net, c, &vec![0; c]).unwrap();
        let v = wai::int_ops::run1(&g, wai_ops::Tensor::new(c as u32, h as u32, w as u32, ints(&zs["z"])), &[]).unwrap();
        assert_eq!(v.data, want, "{plane}: the scale graph's I_s");
        println!("  hyper-scale scale graph (wai-ops, acc32) [{plane}]: {} values bit-exact", v.data.len());
    }
}

/// Close the sigma loop on real data through the integer operator set
/// (integer-operators §13.6.1 steps 2 and 3): z → the plane's scale graph with its gains
/// → `block_stat`, `band_lookup` with the GRFS flags the picture header signals,
/// `scale_index` and `coded`. The RVS additions must be the reference's, and every σ
/// index and mask the one every real `decode_sgm` call consumed.
#[test]
fn jpegai_sigma_chain_end_to_end() {
    use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, Marker};
    use wai_ops::stage;
    let assets = std::path::Path::new("../jpegai-vm-assets");
    let paths = [
        assets.join("hyper_scale_decoder_int_weights.json"),
        assets.join("e2e/hyper_scale_io_fixture.json"),
        assets.join("e2e/sigma_additions_fixture.json"),
        assets.join("e2e/jpegai_real_entropy_fixture.json"),
    ];
    if paths.iter().any(|p| !p.exists()) {
        eprintln!("skipping: sigma-chain fixtures not present");
        return;
    }
    let load = |p: &std::path::Path| -> serde_json::Value {
        serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap()
    };
    let (weights, io, adds, ent) =
        (load(&paths[0]), load(&paths[1]), load(&paths[2]), load(&paths[3]));
    let ints = |x: &serde_json::Value| -> Vec<i64> {
        x.as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect()
    };
    let bin = std::fs::read(assets.join("e2e/test256.bin")).unwrap();
    let segs = parse_container(&bin);
    let ph = parse_picture_header(&bin, find_segment(&segs, Marker::Pih).unwrap().start_bit());
    let (coded_above, edges, additions) = wai::jpegai_ops::stage_constants();
    let table: Vec<Vec<i64>> = additions.iter().map(|r| r.to_vec()).collect();
    for (k, (plane, stream)) in [("y", "SORP"), ("uv", "SORS")].into_iter().enumerate() {
        let shape: Vec<usize> = io[plane]["z_shape"].as_array().unwrap().iter()
            .map(|n| n.as_u64().unwrap() as usize).collect();
        let gain = ints(&adds[plane]["gain_per_ch"]);
        let rvs = ints(&adds[plane]["rvs"]);
        let c = gain.len();
        let grfs = &ph.components[k].grfs_channel_flags;
        let g = wai::jpegai_ops::load_scale(wai::jpegai_ops::scale_net_of_blocks(&weights[plane]).unwrap(), c, &gain).unwrap();
        let v = wai::int_ops::run1(&g, wai_ops::Tensor::new(c as u32, shape[2] as u32, shape[3] as u32, ints(&io[plane]["z"])), &[]).unwrap();
        let hw = v.data.len() / c;
        let (b, sh, lo, hi) = wai::jpegai_ops::LIKELY;
        let likely = stage::block_stat(&v, b, sh, lo, hi).unwrap().data;
        let calls = ent[stream]["calls"].as_array().unwrap();
        assert_eq!(calls.len(), c, "{plane}: call count vs channels");
        for (i, (&x, &l)) in v.data.iter().zip(&likely).enumerate() {
            let row = if grfs[i / hw] { 0 } else { 1 };
            let r = stage::band_lookup(l, &edges, &table, row).unwrap();
            assert_eq!(r, rvs[i], "{plane}: E7 rvs at {i}");
            let call = &calls[i / hw];
            assert_eq!(stage::scale_index(x + r, 7, 32).unwrap(), call["indexes"][i % hw].as_i64().unwrap(), "{plane}: E5 sigma index at {i}");
            assert_eq!(stage::coded(x, coded_above), call["masks"][i % hw].as_i64().unwrap() != 0, "{plane}: E4 mask at {i}");
        }
        println!("  sigma chain through wai-ops [{plane}]: rvs, sigma indexes and masks of {c} channels reproduced from z");
    }
}

/// EFElinear post-filter: Rust implementation vs the live reference filter
/// on the real decode (effective signalled coefficients dumped mid-decode).
#[test]
fn jpegai_efe_linear_matches_reference() {
    use wai::neural::jpegai_filters::{efe_linear_plane, EfeLinearPlane};
    let fx_path = std::path::Path::new("../jpegai-vm-assets/e2e/efl_fixture.json");
    if !fx_path.exists() {
        eprintln!("skipping: EFElinear fixture not present");
        return;
    }
    let fx: serde_json::Value =
        serde_json::from_slice(&std::fs::read(fx_path).unwrap()).unwrap();
    let f32s = |x: &serde_json::Value| -> Vec<f32> {
        x.as_array().unwrap().iter().map(|n| n.as_f64().unwrap() as f32).collect()
    };
    let arr4 = |x: &serde_json::Value| -> [f32; 4] {
        let v = f32s(x);
        [v[0], v[1], v[2], v[3]]
    };
    let input = f32s(&fx["input"]);
    let want = f32s(&fx["want"]);
    let (h, w) = (256usize, 256usize);
    let plane = |i: usize| &input[i * h * w..(i + 1) * h * w];
    // Y must be untouched
    assert_eq!(plane(0), &want[0..h * w], "EFElinear must not modify Y");
    let mut worst = 0f32;
    for (ci, pj) in [(1usize, 0usize), (2, 1)] {
        let pv = &fx["planes"][pj];
        let p = EfeLinearPlane {
            mean: pv["mean"].as_f64().unwrap() as f32,
            w_uv: arr4(&pv["w_uv"]),
            bias: arr4(&pv["bias"]),
            w_y: arr4(&pv["w_y"]),
        };
        let got = efe_linear_plane(plane(0), plane(ci), h, w, &p);
        for (g, wv) in got.iter().zip(&want[ci * h * w..(ci + 1) * h * w]) {
            worst = worst.max((g - wv).abs());
        }
    }
    assert!(worst < 1e-3, "EFElinear worst error {worst}");
    println!("  EFElinear: both chroma planes reproduced, worst |err| = {worst:.2e}");
}

/// LEF post-filter: Rust vs the live reference on the real decode.
#[test]
fn jpegai_lef_matches_reference() {
    use wai::neural::jpegai_filters::lef_apply;
    let fx_path = std::path::Path::new("../jpegai-vm-assets/e2e/lef_fixture.json");
    if !fx_path.exists() {
        eprintln!("skipping: LEF fixture not present");
        return;
    }
    let fx: serde_json::Value =
        serde_json::from_slice(&std::fs::read(fx_path).unwrap()).unwrap();
    let f32s = |k: &str| -> Vec<f32> {
        fx[k].as_array().unwrap().iter().map(|n| n.as_f64().unwrap() as f32).collect()
    };
    let sigma: Vec<i64> = fx["sigma_16"].as_array().unwrap().iter()
        .map(|n| n.as_i64().unwrap()).collect();
    let got = lef_apply(&f32s("y_in"), &sigma, 256, 256, [600, 1200, 2000], [1.1, 1.13, 1.16]);
    let want = f32s("y_want");
    let worst = got.iter().zip(&want).map(|(g, w)| (g - w).abs()).fold(0f32, f32::max);
    assert!(worst < 1e-3, "LEF worst error {worst}");
    println!("  LEF: luma sharpening reproduced, worst |err| = {worst:.2e}");
}

/// EFEnonlinear post-filter: Rust vs the live reference on the real decode.
#[test]
fn jpegai_efe_nonlinear_matches_reference() {
    use wai::neural::jpegai_filters::efe_nonlinear_plane;
    let fx_path = std::path::Path::new("../jpegai-vm-assets/e2e/efn_fixture.json");
    if !fx_path.exists() {
        eprintln!("skipping: EFEnonlinear fixture not present");
        return;
    }
    let fx: serde_json::Value =
        serde_json::from_slice(&std::fs::read(fx_path).unwrap()).unwrap();
    let f32s = |k: &str| -> Vec<f32> {
        fx[k].as_array().unwrap().iter().map(|n| n.as_f64().unwrap() as f32).collect()
    };
    let wq: Vec<i64> = fx["weights_q15"].as_array().unwrap().iter()
        .map(|n| n.as_i64().unwrap()).collect();
    let got = efe_nonlinear_plane(
        &f32s("y"), &f32s("c_in"),
        &[wq[0], wq[1], wq[2], wq[3], wq[4], wq[5], wq[6], wq[7]],
        fx["luma_min_q"].as_i64().unwrap(), fx["luma_max_q"].as_i64().unwrap());
    let want = f32s("c_want");
    let worst = got.iter().zip(&want).map(|(g, w)| (g - w).abs()).fold(0f32, f32::max);
    assert!(worst < 1e-3, "EFEnonlinear worst error {worst}");
    println!("  EFEnonlinear: V plane reproduced, worst |err| = {worst:.2e}");
}

/// Header parser vs the live-decode field trace: every typed field must
/// match what the reference decoder read from PIH + TON.
#[test]
fn jpegai_header_parser_matches_trace() {
    use wai::neural::jpegai_header::parse_headers;
    let assets = std::path::Path::new("../jpegai-vm-assets");
    let bin_path = assets.join("e2e/test256.bin");
    if !bin_path.exists() {
        eprintln!("skipping: bitstream not present");
        return;
    }
    let bin = std::fs::read(&bin_path).unwrap();
    let subs = wai::neural::jpegai_bitstream::parse_codestream(&bin).unwrap();
    let payload = |m: u16| subs.iter().find(|s| s.marker == m).unwrap().payload;
    let h = parse_headers(payload(0xFF82), payload(0xFF83)).expect("parse");
    assert_eq!(h.stream_profile_idc, 0);
    assert_eq!(h.decoder_profile_id, 2);
    assert_eq!(h.synthesis_transform_ids, vec![2, 1, 0]);
    assert_eq!(h.level_idc, 52);
    assert_eq!((h.img_width, h.img_height), (256, 256));
    assert_eq!(h.colour_transform_idx, 1);
    assert_eq!(h.model_id, 0);
    assert_eq!(h.beta_displacement_log, 0);
    assert_eq!(h.models[0].num_chs, 160);
    assert_eq!(h.models[1].num_chs, 96);
    assert_eq!(h.models[0].grfs_channel_flag.len(), 160);
    assert_eq!(h.models[1].grfs_channel_flag.len(), 96);
    assert_eq!(&h.models[0].grfs_channel_flag[..16],
               &[1, 1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0]);
    assert!(h.efe_linear_enabled);
    assert_eq!((h.efe_linear.mean_u_q100, h.efe_linear.mean_v_q100), (12561, 13091));
    assert_eq!(h.efe_linear.weights_u, vec![32677]);
    assert_eq!(h.efe_linear.weights_v, vec![32758]);
    assert_eq!(h.efe_linear.weights_u2, vec![32772]);
    assert_eq!(h.efe_linear.weights_v2, vec![32776]);
    assert!(h.icci_enabled && h.icci_short_list);
    assert_eq!(h.icci_signalled_idx, [0, 1]);
    assert!(h.efe_nonlinear_enabled && !h.efn_u_enabled && h.efn_v_enabled);
    assert_eq!((h.efn_min_luma_q100, h.efn_max_luma_q100), (497, 25320));
    assert_eq!(h.efn_weights, vec![32966, 32443, 33031, 32361, 33068, 33006, 31968, 33003]);
    assert!(h.lef_enabled);
    assert_eq!(h.lef_ch_idx, 0);
    println!("  header parser: all typed fields match the live-decode trace");
}


/// The residual-variance bands, against the reference decoder's own tables: for every
/// `likely` (0 … 4351) and both GRFS branches, the band the exact decoder's edges give
/// selects exactly the σ addition and the dequant multiplier the reference's 4352-entry
/// tables hold at that index (`model_constants.json`, exported from the reference
/// decoder). The tables change value at 1004, 1353 and 1938.
#[test]
fn jpegai_rvs_bands_are_the_reference_tables() {
    use wai::jpegai_pipeline::rvs_tables;
    let (_, edges, additive) = wai::jpegai_ops::stage_constants();
    let path = std::path::Path::new("../jpegai-vm-assets/model_constants.json");
    if !path.exists() {
        eprintln!("skipping: model constants not present");
        return;
    }
    let mc: serde_json::Value = serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap();
    let ints = |k: &str| -> Vec<i64> {
        mc[k].as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect()
    };
    for flag in [false, true] {
        let add = ints(&format!("rvs_sigma_t_flag{}", flag as u8));
        let dq = ints(&format!("rvs_dq_t_flag{}", flag as u8));
        assert_eq!((add.len(), dq.len()), (4352, 4352));
        let (_, dq_table, add_table) = rvs_tables(flag);
        assert_eq!(add_table, additive[if flag { 0 } else { 1 }], "the additive table is its derivation");
        for likely in 0..4352i64 {
            let band = edges.iter().filter(|&&e| likely >= e).count();
            let i = likely as usize;
            assert_eq!(add_table[band], add[i], "σ addition, flag {flag}, likely {likely}");
            assert_eq!(dq_table[band], dq[i], "dequant multiplier, flag {flag}, likely {likely}");
        }
    }
    println!("  RVS bands: both branches' σ additions and dequant multipliers equal the reference's tables at all 4352 indices");
}