open-standards/wai/wai-rs/tests/jpegai_bitstream_calibration.rs

//! Calibration for [`wai::jpegai_bitstream`] against a REAL JPEG-AI codestream
//! (`test256.bin`): walk the marker segments and check the layout, then replay the
//! picture header's leading fields with the bit reader and match the reference values.
//!
//! The strongest check here is a cross-validation: the SORP / SORS / SOZ payload sizes
//! recovered by the container parse must equal the `mem` lengths in the entropy fixture
//! that the byte-exact mANS decoder consumes — two independent routes to the same bytes.
//!
//! env: WAI_JPEGAI_CODESTREAM = path to test256.bin (optional; defaults to the committed
//!      copy in tests/vectors/jpegai/)
#![cfg(feature = "neural_int")]

use wai::jpegai_bitstream::{find_segment, parse_container, parse_lsbs_flags, parse_picture_header,
                            parse_tools_header, BitReader, Marker};

/// `WAI_JPEGAI_CODESTREAM` if set, else the committed copy of the same reference
/// codestream (`tests/vectors/jpegai/test256.bin`), so these tests do not skip.
fn codestream() -> Option<Vec<u8>> {
    match std::env::var("WAI_JPEGAI_CODESTREAM") {
        Ok(p) => Some(std::fs::read(p).expect("WAI_JPEGAI_CODESTREAM is set but unreadable")),
        Err(_) => Some(include_bytes!("vectors/jpegai/test256.bin").to_vec()),
    }
}

#[test]
fn container_layout_matches_reference_codestream() {
    let Some(data) = codestream() else { return };
    let segs = parse_container(&data);
    for s in &segs {
        eprintln!("  {:?} payload bytes {}..{} (len {})", s.marker, s.start, s.start + s.len, s.len);
    }
    let order: Vec<Marker> = segs.iter().map(|s| s.marker).collect();
    assert_eq!(
        order,
        vec![Marker::Soc, Marker::Pih, Marker::Ton, Marker::Rdi, Marker::Sorp, Marker::Sors, Marker::Soz, Marker::Eoc],
        "codestream segment order"
    );
    let pih = find_segment(&segs, Marker::Pih).unwrap();
    assert_eq!((pih.start, pih.len), (6, 49), "PIH payload");
    let ton = find_segment(&segs, Marker::Ton).unwrap();
    assert_eq!((ton.start, ton.len), (59, 63), "TON payload");

    // Cross-validation: the entropy substream sizes must equal the mANS fixture `mem`
    // lengths (SORP=97, SORS=47, SOZ=83) — the container parse and the byte-exact
    // entropy decode agree on the same bytes from opposite directions.
    assert_eq!(find_segment(&segs, Marker::Sorp).unwrap().len, 97, "SORP == mANS mem");
    assert_eq!(find_segment(&segs, Marker::Sors).unwrap().len, 47, "SORS == mANS mem");
    assert_eq!(find_segment(&segs, Marker::Soz).unwrap().len, 83, "SOZ == mANS mem");
}

#[test]
fn picture_header_leading_fields_match_reference() {
    let Some(data) = codestream() else { return };
    let segs = parse_container(&data);
    let pih = find_segment(&segs, Marker::Pih).expect("PIH present");
    let mut br = BitReader::at(&data, pih.start_bit());
    assert_eq!(br.bit_pos(), 48, "PIH payload starts at bit 48");

    // The picture header's fixed prefix, per the reference syntax.
    let stream_profile_idc = br.read(4);
    let decoder_profile_id = br.read(4);
    let num_synthesis_transforms_minus1 = br.read(4);
    let synthesis_transform_id: Vec<u64> = (0..=num_synthesis_transforms_minus1).map(|_| br.read(4)).collect();
    let level_idc = br.read(8);
    let img_width = br.read(16) + 64;
    let img_height = br.read(16) + 64;
    let _diff_display_w = br.read(6);
    let _diff_display_h = br.read(6);
    let bit_depth_idc = br.read(3);
    let (s_ver, s_hor) = (br.read(1) + 1, br.read(1) + 1);
    let (c_ver, c_hor) = (br.read(1) + 1, br.read(1) + 1);
    let colour_transform_idx = br.read(2);
    let model_id = br.read(4);

    eprintln!("profile={stream_profile_idc} decoder={decoder_profile_id} level={level_idc} \
               {img_width}x{img_height} bd_idc={bit_depth_idc} colour_idx={colour_transform_idx} model={model_id}");

    assert_eq!(stream_profile_idc, 0);
    assert_eq!(decoder_profile_id, 2);
    assert_eq!(num_synthesis_transforms_minus1, 2);
    assert_eq!(synthesis_transform_id, vec![2, 1, 0]);
    assert_eq!(level_idc, 52);
    assert_eq!((img_width, img_height), (256, 256), "test256 is 256x256");
    assert_eq!(bit_depth_idc, 0);
    assert_eq!((s_ver, s_hor, c_ver, c_hor), (1, 1, 1, 1), "4:4:4");
    assert_eq!(colour_transform_idx, 1, "built-in YUV<->RGB");
    assert_eq!(model_id, 0);
    assert_eq!(br.bit_pos(), 137, "cursor lands where the reference trace does");
}

#[test]
fn full_picture_header_matches_reference() {
    let Some(data) = codestream() else { return };
    let segs = parse_container(&data);
    let pih = find_segment(&segs, Marker::Pih).expect("PIH present");
    let h = parse_picture_header(&data, pih.start_bit());
    eprintln!("{h:#?}");

    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.bit_depth_idc, 0);
    assert_eq!((h.s_ver, h.s_hor, h.c_ver, h.c_hor), (1, 1, 1, 1), "4:4:4");
    assert_eq!(h.colour_transform_idx, 1);
    assert!(h.colour_transform_matrix.is_none(), "idx=1 carries no matrix");
    assert_eq!(h.model_id, 0);
    assert_eq!(h.num_threads_z, 1);
    assert_eq!(h.beta_displacement_log, [0, 0], "2048 de-biased, shared across Y/UV");
    assert!(!h.independent_beta_uv);
    assert!(!h.region_partitioning);
    assert!(!h.gain_3d_enable);

    // per-component (luma 160ch, chroma 96ch); the four PIH tool flags as decoded
    assert_eq!(h.components[0].num_chs, 160);
    assert_eq!(h.components[1].num_chs, 96);
    for c in 0..2 {
        let comp = &h.components[c];
        assert!(!comp.use_cube_flags, "c{c} use_cube");
        assert!(comp.rvs_enable && comp.grfs_enable, "c{c} rvs/grfs precede the channel flags");
        assert_eq!(comp.grfs_channel_flags.len(), comp.num_chs as usize, "c{c} grfs flag count == num_chs");
        assert!(!comp.synthesis_tile_enable, "c{c} synthesis tile flag trails the channel flags");
    }
    // LSBS codes its flags into the tool header, not here: the first two TON bits,
    // both 1 (reference trace: lsbs_enable_flag[0] = lsbs_enable_flag[1] = 1).
    let ton = find_segment(&segs, Marker::Ton).expect("TON present");
    assert_eq!(parse_lsbs_flags(&data, ton.start_bit()), [true, true], "LSBS on for Y and UV");
    // GRFS channel-selection counts — independently confirmed by deriving the same
    // per-channel cwg assignment from the reference RVS output (both planes byte-exact)
    assert_eq!(h.components[0].grfs_channel_flags.iter().filter(|&&b| b).count(), 5);
    assert_eq!(h.components[1].grfs_channel_flags.iter().filter(|&&b| b).count(), 10);

    // the whole header consumed, leaving only the zero padding of the 49-byte PIH
    let pih_end = (pih.start + pih.len) * 8;
    assert_eq!(h.end_bit, 435, "all PIH fields consumed");
    assert!(pih_end - h.end_bit < 8, "only byte-alignment padding remains");
    let mut br = BitReader::at(&data, h.end_bit);
    assert_eq!(br.read((pih_end - h.end_bit) as u32), 0, "the padding bits are zero");
}

#[test]
fn tools_header_matches_reference() {
    let Some(data) = codestream() else { return };
    let segs = parse_container(&data);
    let ton = find_segment(&segs, Marker::Ton).expect("TON present");
    let pih = parse_picture_header(&data, find_segment(&segs, Marker::Pih).unwrap().start_bit());

    // 4:4:4 ⇒ scale_ver = scale_hor = 3 − c/s = 2 ⇒ the single-plane filter path
    let scale_2x2 = (3 - pih.c_ver / pih.s_ver == 2) && (3 - pih.c_hor / pih.s_hor == 2);
    assert!(scale_2x2, "test256 is 4:4:4");

    // The post-filter chain starts 2 bits into TON: `_stream_header_part` defaults to
    // "tool_header", so other tools emit into TON ahead of the filters. models_count
    // sets the signalled-index width (2 bits here).
    let filters_bit = ton.start_bit() + 2;
    let t = parse_tools_header(&data, filters_bit, scale_2x2, 3, 1);
    eprintln!("{t:#?}");

    // EFE linear: the two plane means feeding efl_chroma
    assert!(t.efl_enabled);
    let efl = t.efl.as_ref().unwrap();
    assert_eq!(efl.b1, [12561, 13091], "B1/100 = 125.61 / 130.91 = the EFL plane means");
    let f1 = &efl.filters_best2;
    assert_eq!((f1.cand_u, f1.cand_v), (0, 0));
    assert_eq!((f1.fl_u, f1.fl_v), (1, 1));
    assert_eq!((f1.min_symbol, f1.max_symbol), (0, 65535));
    assert_eq!((&f1.weights_u[..], &f1.weights_u2[..]), (&[32677u32][..], &[32772u32][..]));
    assert_eq!((&f1.weights_v[..], &f1.weights_v2[..]), (&[32758u32][..], &[32776u32][..]));
    let f2 = &efl.filters_best;
    // the second group's max_symbol (99) narrows the weights to 7 bits
    assert_eq!((f2.min_symbol, f2.max_symbol), (32677, 99));
    assert_eq!((&f2.weights_u[..], &f2.weights_u2[..]), (&[0u32][..], &[95u32][..]));
    assert_eq!((&f2.weights_v[..], &f2.weights_v2[..]), (&[81u32][..], &[99u32][..]));

    // eICCI: the per-tile model selection (the last input the eICCI net needed)
    assert!(t.icci_enabled);
    assert!(!t.icci_tile_manager_enabled);
    assert_eq!(t.icci_tiles.len(), 1);
    let tile = &t.icci_tiles[0];
    assert_eq!(tile.use_yuv, [true, true, true]);
    assert!(tile.use_short_list);
    assert_eq!((tile.model_idx_y, tile.model_idx_uv), (0, 1));

    // EFE nonlinear: exactly the inputs efn_filter consumes
    assert!(t.efn_enabled);
    let efn = t.efn.as_ref().unwrap();
    assert_eq!((efn.min_symbol, efn.max_symbol), (0, 65535));
    assert!(!efn.mask1_enabled && !efn.mask2_enabled);
    assert!(!efn.nonlinear_u_enabled && efn.nonlinear_v_enabled, "only V is filtered");
    assert_eq!((efn.tile_width, efn.tile_height), (256, 256));
    assert_eq!((&efn.luma_min[..], &efn.luma_max[..]), (&[497u32][..], &[25320u32][..]),
               "== efn_fixture luma_min_q / luma_max_q");
    assert!(efn.weights_u.is_empty(), "U disabled ⇒ no candNum/A1 for U");
    assert_eq!(efn.weights_v, vec![32966, 32443, 33031, 32361, 33068, 33006, 31968, 33003],
               "== efn_fixture weights_q15");

    assert!(t.lef_enabled);
    assert_eq!(t.lef_ch_idx, 0);

    // the whole tools header consumed exactly, up to the pad closing the 63-byte TON
    assert_eq!(t.end_bit, 973, "all TON fields consumed");
    assert!(t.end_bit <= (ton.start + ton.len) * 8);
}

#[test]
fn header_widths_follow_calc_bits_count() {
    use wai::jpegai_bitstream::bits_for_max;
    // bits = ceil(log2(max+1)) — the reference HeaderCoder::calc_bits_count
    for (max, want) in [(0u64, 0u32), (1, 1), (2, 2), (3, 2), (4, 3), (9, 4), (15, 4), (99, 7), (255, 8), (65535, 16)] {
        assert_eq!(bits_for_max(max), want, "bits_for_max({max})");
    }
}