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

//! End-to-end **chroma**: raw codestream → the decoded U and V planes, no per-image
//! fixture — the counterpart to `jpegai_e2e_frontend`'s luma chain.
//!
//! ```text
//! test256.bin → SOZ → mANS fac (luma then uv CDFs) → z_uv
//!             → decode_hyper_scale                 → i_s
//!             → rvs_likely / rvs_additive / skip   → sigma, mask
//!             → SORS → mANS sgm                    → y_q      bit-exact
//!             → dequantize_resi                    → resi     bit-exact
//!             → decode_hyper                       → ψ
//!             → reconstruct_yhat_uv(resi, ψ)       → y_hat
//!             → lsbs_refine                        → uv_lsbs
//!             → decode_uv(cat(y_lsbs, uv_lsbs))    → U,V[256²]
//! ```
//!
//! Chroma is **not** luma with different constants — two structural differences make
//! this worth exercising separately rather than parameterising the luma test:
//!
//!   * **No MCM.** The reference ships no `MCM/` nets under `model_uv`; the hyper
//!     params are used directly as the means (see [`reconstruct_yhat_uv`]), where luma
//!     runs four autoregressive context stages.
//!   * **Synthesis is joint.** `decode_uv` consumes the *concatenated* 256-channel
//!     latent `cat(y_lsbs, uv_lsbs)`, so chroma pixels depend on the luma chain being
//!     right as well — a luma regression surfaces here too.
//!
//! One SOZ substream carries both planes' hyper-latents back to back through a single
//! decoder whose ANS state threads across them, so the luma pass is decoded first even
//! though only the chroma output is checked here.
//!
//! Everything upstream of synthesis is shared code with luma, which is the point: the
//! same derivations (`rvs_likely`, `skip_mask`, `sigma_index`, `dequantize_resi`) are
//! re-exercised on a second plane with different channel counts and flags.
//!
//! env: WAI_JPEGAI_FRONTEND_FIXTURE (uses its `uv` sub-object and the luma `y_lsbs` path)
//!      WAI_JPEGAI_HYPERSCALE_UV_WEIGHTS, WAI_JPEGAI_HYPER_UV_WEIGHTS,
//!      WAI_JPEGAI_SYNTH_UV_WEIGHTS  (optional; each stage is skipped without its weights)
#![cfg(all(feature = "neural_int", feature = "legacy_int"))]

use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, Marker};
use wai::jpegai_mans::{build_decode_transitions, cdf_from_pmf, MansDecoder};
use wai::jpegai_pipeline::{dequantize_resi, gain_unit_scaler_vec, lsbs_refine, reconstruct_yhat_uv,
    rvs_additive, rvs_likely, sigma_index, skip_mask};
use wai::jpegai_synthesis::{decode_hyper, decode_hyper_scale, q16_from_symbols, IntNet, Weights};

const Z_OFFSET: i64 = 31;

fn env_path(k: &str) -> Option<String> {
    std::env::var(k).ok().filter(|p| std::path::Path::new(p).exists())
}
fn weights(p: &str) -> Weights {
    Weights::from_json(&serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap()).unwrap()
}
fn snr(got_q16: &[i64], expected: &[f32]) -> f64 {
    let (mut se, mut sig) = (0f64, 0f64);
    for (&g, &e) in got_q16.iter().zip(expected) {
        let (q, f) = (g as f64 / 65536.0, e as f64);
        se += (q - f) * (q - f);
        sig += f * f;
    }
    if se == 0.0 { f64::INFINITY } else { 10.0 * (sig / se).log10() }
}

#[test]
fn codestream_to_chroma_pixels() {
    let (Some(fx), Some(hs_uv)) = (env_path("WAI_JPEGAI_FRONTEND_FIXTURE"),
                                   env_path("WAI_JPEGAI_HYPERSCALE_UV_WEIGHTS")) else {
        eprintln!("skip: set WAI_JPEGAI_FRONTEND_FIXTURE + WAI_JPEGAI_HYPERSCALE_UV_WEIGHTS");
        return;
    };
    let root: serde_json::Value = serde_json::from_slice(&std::fs::read(&fx).unwrap()).unwrap();
    let v = &root["uv"];
    let u8v = |o: &serde_json::Value, k: &str| o[k].as_array().unwrap().iter()
        .map(|x| x.as_u64().unwrap() as u8).collect::<Vec<u8>>();
    let i64v = |o: &serde_json::Value, k: &str| o[k].as_array().unwrap().iter()
        .map(|x| x.as_i64().unwrap()).collect::<Vec<i64>>();
    let f32v = |o: &serde_json::Value, k: &str| o[k].as_array().unwrap().iter()
        .map(|x| x.as_f64().unwrap() as f32).collect::<Vec<f32>>();

    let data = u8v(&root, "codestream");
    let (channels, h, w) = (
        v["channels"].as_u64().unwrap() as usize,
        v["h"].as_u64().unwrap() as usize,
        v["w"].as_u64().unwrap() as usize,
    );
    let zs: Vec<usize> = v["z_shape"].as_array().unwrap().iter()
        .map(|x| x.as_u64().unwrap() as usize).collect();
    let (zc, zh, zw) = (zs[1], zs[2], zs[3]);
    let per_channel = root["per_channel"].as_u64().unwrap() as usize;
    let gain_vector_log = i64v(v, "gain_vector_log");

    let segs = parse_container(&data);

    // grfs and beta are DERIVED from the codestream's picture header, not the fixture:
    // component [1] is chroma. This is the same parser `jpegai_header_parser_matches_trace`
    // validates against the reference trace, so nothing per-image is assumed here.
    let pih = find_segment(&segs, Marker::Pih).expect("PIH");
    let hdr = parse_picture_header(&data, pih.start_bit());
    let grfs = hdr.components[1].grfs_channel_flags.clone();
    let beta = hdr.beta_displacement_log[1] as i64;
    assert_eq!(grfs.len(), channels, "one GRFS flag per chroma channel");

    // 1. SOZ → z_uv. One SOZ substream carries BOTH planes' hyper-latents back to
    // back — luma first, then chroma — through a single decoder whose ANS state
    // threads across them. The luma pass is therefore not optional scaffolding here:
    // skipping it leaves the decoder at the wrong stream position and chroma decodes
    // garbage (or underflows outright).
    let soz = find_segment(&segs, Marker::Soz).expect("SOZ");
    let mem = &data[soz.start..soz.start + soz.len];
    let mut dec_z = MansDecoder::new(mem, &[mem.len() as u32]);

    let y_ch = root["z_shape"].as_array().unwrap()[1].as_u64().unwrap() as usize;
    let mut y_syms = vec![0u8; y_ch * per_channel];
    dec_z.decode_factorize(&u8v(&root, "cdfs_y"), &mut y_syms, y_ch, per_channel);

    let mut symbols = vec![0u8; zc * per_channel];
    dec_z.decode_factorize(&u8v(v, "cdfs_uv"), &mut symbols, zc, per_channel);
    let z: Vec<i64> = symbols.iter().map(|&s| s as i64 - Z_OFFSET).collect();

    // 2. z → hyper-scale → i_s, then the derived entropy inputs (shared with luma)
    let net = IntNet::from_json(&serde_json::from_slice(&std::fs::read(&hs_uv).unwrap()).unwrap()).unwrap();
    let (i_s, ..) = decode_hyper_scale(&z, zc, zh, zw, &net);
    assert_eq!(i_s.len(), channels * h * w, "i_s shape");

    let likely = rvs_likely(&i_s, &gain_vector_log, channels, h, w);
    let rvs = rvs_additive(&i_s, &gain_vector_log, &grfs, channels, h, w);
    let sigma = sigma_index(&i_s, &gain_vector_log, &rvs, channels);
    let mask = skip_mask(&i_s, &gain_vector_log, channels);
    eprintln!("[uv] derived: {} coded positions of {}", mask.iter().filter(|&&m| m).count(), mask.len());

    // 3. SORS → mANS sgm → y_q
    let sors = find_segment(&segs, Marker::Sors).expect("SORS");
    let mem = &data[sors.start..sors.start + sors.len];
    let pmf: Vec<Vec<i64>> = root["pmf"].as_array().unwrap().iter()
        .map(|r| r.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect())
        .collect();
    let transitions = build_decode_transitions(&pmf, &cdf_from_pmf(&pmf));
    let flat: Vec<u32> = transitions.iter().flatten().copied().collect();

    let mut dec = MansDecoder::new(mem, &[mem.len() as u32]);
    dec.set_sgm_transitions(&flat, &u8v(&root, "bounds"));
    let plane = h * w;
    let mut y_q = vec![0i16; channels * plane];
    for c in 0..channels {
        let (a, b) = (c * plane, (c + 1) * plane);
        let mut row = vec![0i16; plane];
        dec.decode_sgm(&sigma[a..b], &mut row, &mask[a..b]);
        y_q[a..b].copy_from_slice(&row);
    }
    let y_q_expected: Vec<i16> = i64v(v, "y_q_expected").iter().map(|&x| x as i16).collect();
    let bad = y_q.iter().zip(&y_q_expected).filter(|(a, b)| a != b).count();
    eprintln!("[uv] y_q from codestream: {}/{} exact", y_q.len() - bad, y_q.len());
    assert_eq!(bad, 0, "codestream → y_q (chroma) must be byte-exact");

    // 1–3 again, through the integer operator set as the decoder runs them
    // (integer-operators §13.6.1): a fresh `fac` decode over SOZ that decodes and discards
    // the 160 luma rows (with luma's rows) before chroma's 96, the plane's scale graph
    // loaded under acc32, the stage operators, and a fresh `sgm` decode over SORS. Every
    // integer must be the one above, and y_q the reference's.
    let gain_uv = &gain_vector_log;
    let graph = wai::jpegai_ops::load_scale(wai::jpegai_ops::scale_net(&net).unwrap(), zc, gain_uv).expect("hs_uv loads under acc32");
    let rows = wai::jpegai_ops::fac_rows(&i64v(&root, "cdfs_y"), &i64v(v, "cdfs_uv"), wai::jpegai_ops::Plane::Chroma);
    let tables = wai_ops::stage::mans::build_latent_tables(&pmf).unwrap();
    let bounds = i64v(&root, "bounds");
    let setup = wai::jpegai_ops::PlaneSetup { graph: &graph, rows: &rows, skip: y_ch, tables: &tables, bounds: &bounds };
    let ops = wai::jpegai_ops::plane_with(&setup, &data[soz.start..soz.start + soz.len], mem, &grfs).expect("the chroma plane decodes through the operator set");
    assert_eq!(ops.z, z, "[uv] z through the operator set");
    let v_direct: Vec<i64> = i_s.iter().enumerate().map(|(i, &x)| x + gain_uv[i / (h * w)]).collect();
    assert_eq!(ops.v, v_direct, "[uv] v through the operator set");
    assert_eq!((&ops.likely, &ops.rvs, &ops.mask), (&likely, &rvs, &mask), "[uv] likely, rvs and mask through the operator set");
    assert_eq!(ops.sigma, sigma.iter().map(|&s| s as i64).collect::<Vec<_>>(), "[uv] σ_idx through the operator set");
    let ops_bad = ops.y_q.iter().zip(&y_q_expected).filter(|(a, b)| **a != **b as i64).count();
    eprintln!("[uv] y_q through the operator set: {}/{} exact", ops.y_q.len() - ops_bad, ops.y_q.len());
    assert_eq!(ops_bad, 0, "codestream → y_q (chroma) through the operator set must be byte-exact");

    // 4. dequant — bit-exact against the reference residual
    let sv = gain_unit_scaler_vec(&gain_vector_log, beta);
    let resi = dequantize_resi(&y_q, &likely, &grfs, &sv, channels);
    let resi_exp = f32v(v, "resi_full_expected");
    let exact = resi.iter().zip(&resi_exp).filter(|(a, b)| a == b).count();
    eprintln!("[uv] resi from codestream: {exact}/{} bit-exact", resi_exp.len());
    assert_eq!(exact, resi_exp.len(), "chroma dequant must be bit-exact");

    // 5. ψ → reconstruction (no MCM) → LSBS
    let Some(hyp_uv) = env_path("WAI_JPEGAI_HYPER_UV_WEIGHTS") else {
        eprintln!("[uv] skip ψ onward: set WAI_JPEGAI_HYPER_UV_WEIGHTS");
        return;
    };
    let (psi, ..) = decode_hyper(&q16_from_symbols(&z), zc, zh, zw, &weights(&hyp_uv));
    eprintln!("[uv] ψ: SNR {:.2} dB", snr(&psi, &f32v(v, "psi_expected")));

    let resi_q16: Vec<i64> = resi.iter().map(|&r| (r as f64 * 65536.0).round() as i64).collect();
    let yhat = reconstruct_yhat_uv(&resi_q16, &psi, channels, h, w);
    let yhat_snr = snr(&yhat, &f32v(v, "yhat_expected"));
    eprintln!("[uv] y_hat: SNR {yhat_snr:.2} dB");
    assert!(yhat_snr > 40.0, "chroma y_hat SNR {yhat_snr:.2} dB");

    let t0: Vec<i32> = root["lsbs_t0"].as_array().unwrap().iter()
        .map(|x| x.as_i64().unwrap() as i32).collect();
    let t1: Vec<i32> = root["lsbs_t1"].as_array().unwrap().iter()
        .map(|x| x.as_i64().unwrap() as i32).collect();
    let yhat_f: Vec<f32> = yhat.iter().map(|&q| q as f64 as f32 / 65536.0).collect();
    let uv_lsbs = lsbs_refine(&yhat_f, &resi, &likely, &t0, &t1, channels, h, w);
    let lsbs_exp = f32v(v, "lsbs_out_expected");
    let (mut se, mut sig, mut id_se) = (0f64, 0f64, 0f64);
    for i in 0..lsbs_exp.len() {
        let (g, e, id) = (uv_lsbs[i] as f64, lsbs_exp[i] as f64, yhat_f[i] as f64);
        se += (g - e) * (g - e);
        sig += e * e;
        id_se += (id - e) * (id - e);
    }
    let (lsnr, id_snr) = (10.0 * (sig / se).log10(), 10.0 * (sig / id_se).log10());
    eprintln!("[uv] uv_lsbs: SNR {lsnr:.2} dB (identity baseline {id_snr:.2} dB)");
    assert!(lsnr > id_snr + 10.0, "chroma LSBS no better than identity");

    // 6. Joint synthesis: cat(y_lsbs, uv_lsbs) → U,V. Needs the luma latent, so this
    // also re-checks the luma chain end to end.
    let (Some(sw_uv), Some(y_lsbs_j)) = (env_path("WAI_JPEGAI_SYNTH_UV_WEIGHTS"),
                                         root.get("lsbs_out_expected")) else {
        eprintln!("[uv] skip synthesis: set WAI_JPEGAI_SYNTH_UV_WEIGHTS");
        return;
    };
    let mut lat: Vec<i64> = y_lsbs_j.as_array().unwrap().iter()
        .map(|x| (x.as_f64().unwrap() * 65536.0).round() as i64).collect();
    lat.extend(uv_lsbs.iter().map(|&f| (f as f64 * 65536.0).round() as i64));
    let uv_pix = wai::jpegai_synthesis::decode_uv(&lat, &weights(&sw_uv));
    let exp = f32v(v, "synth_uv_expected");
    assert_eq!(uv_pix.len(), exp.len(), "chroma plane size");
    let (mut maxd, mut off) = (0i32, 0usize);
    for (&g, &e) in uv_pix.iter().zip(&exp) {
        let d = (g as i32 - e.round().clamp(0.0, 255.0) as i32).abs();
        maxd = maxd.max(d);
        if d != 0 { off += 1; }
    }
    let total = exp.len();
    eprintln!("[uv] U,V planes: max|Δ| = {maxd}, {off}/{total} differ ({:.3}%)",
              100.0 * off as f64 / total as f64);

    // Attribute the pixel error rather than just widening the bar: rerun synthesis on
    // the REFERENCE latent for both planes. What remains is the synthesis net's own
    // 16-bit PTQ; the difference is everything the chain accumulated upstream.
    let mut ref_lat: Vec<i64> = y_lsbs_j.as_array().unwrap().iter()
        .map(|x| (x.as_f64().unwrap() * 65536.0).round() as i64).collect();
    ref_lat.extend(f32v(v, "lsbs_out_expected").iter().map(|&f| (f as f64 * 65536.0).round() as i64));
    let uv_ref = wai::jpegai_synthesis::decode_uv(&ref_lat, &weights(&sw_uv));
    let roff = uv_ref.iter().zip(&exp)
        .filter(|(g, e)| **g as i32 != e.round().clamp(0.0, 255.0) as i32).count();
    eprintln!("[uv] from REFERENCE latent: {roff}/{total} differ ({:.3}%)  \
               ⇒ synthesis-port {:.3}%, accumulated upstream {:.3}%",
              100.0 * roff as f64 / total as f64,
              100.0 * roff as f64 / total as f64,
              100.0 * (off as f64 - roff as f64) / total as f64);

    assert!(maxd <= 1, "chroma deviation beyond rounding: {maxd}");
    // Chroma runs looser than luma (~3.4% vs ~1.4% of samples at 1 LSB). Its ψ is
    // 60.6 dB where luma's is 70.5, and with no MCM to refine them those hyper params
    // are used directly as the means, so the quantization error reaches the latent
    // unattenuated. The bar is set from the measured value with headroom, and the
    // attribution above is what a regression would actually move.
    assert!(off * 100 <= total * 4, "more than 4% mismatched: {off}/{total}");
}