//! The integer stages of `wai.image.jpegai` through the integer operator set
//! `wai-ops/1`, lowered by `jpegai-int/1`: the scale path and the two entropy decodes of
//! `extensions/integer-operators.md` §13.6.1, and the hyper decoders, context stages and
//! synthesis of §13.6.2 ([`Networks`]).
//!
//! A bundle's scale path is lowered and loaded once ([`load`], [`load_bytes`]): the
//! load checks of §13.6.1 on `norm.json`, the two scale graphs `hs_y` and `hs_uv`
//! loaded under the capability's profile ([`PROFILE`]), and the latent tables E8 builds
//! from `pmf`. A [`ScalePath`] is what a decode then holds of the set. [`plane`] decodes
//! one plane's integers from the codestream's entropy substreams, as §13.6.1 steps 1–4
//! give them: the hyper-latent (`fac`), the scale graph, `coded`, `block_stat`,
//! `band_lookup`, `scale_index`, and the latent (`sgm`).
//!
//! Every refusal is the operator set's: `model_mismatch` for a bundle that fails a load
//! check or whose graphs do not meet the profile, and `stream_exhausted` for a stream
//! the entropy decode cannot read (E8). Nothing here panics on a stream.

use wai_ops::lower::jpegai::{JpegaiNet, Member, Norm, SCALE_MEMBERS, ScaleNet};
use wai_ops::stage::{self, mans};
use wai_ops::{Code, Loaded, Profile, Refusal, Site, Tensor};

use crate::jpegai_decode::NormTables;
use crate::jpegai_pipeline::{RVS_ADDITIVE, RVS_BAND_EDGES};
use crate::jpegai_synthesis::IntNet;

/// The lowering the capability's scale path is lowered by (§13.6.1), as the registry's
/// `ops` member names it.
pub const LOWERING: &str = "jpegai-int/1";

/// The profile `wai.image.jpegai` declares (§11): every conforming bundle's scale
/// graphs meet `acc32`, so they run on a 32-bit fabric as on a 64-bit one. A bundle
/// whose certificate does not meet it is refused (`model_mismatch`).
pub(crate) const PROFILE: Profile = Profile::Acc32;

/// The hyper-latent symbols' offset: `z = sym − 31` (§13.6.1 step 1).
pub(crate) const Z_OFFSET: i64 = 31;
/// The `coded` threshold: a position is coded when `v > 382` (step 3).
pub(crate) const CODED_ABOVE: i64 = 382;
/// The parameters of step 3's stage operators, for a caller outside the crate: the
/// `coded` threshold, then `block_stat`'s and `scale_index`'s (below), and the RVS band
/// edges and additions [`band_lookup`](wai_ops::stage::band_lookup) reads. (The integer
/// constants are crate-only, so that the C header does not carry them.)
pub fn stage_constants() -> (i64, [i64; 3], [[i64; 4]; 2]) {
    (CODED_ABOVE, RVS_BAND_EDGES, RVS_ADDITIVE)
}
/// `block_stat(v; 8, 6, 0, 4351)` (step 3).
pub const LIKELY: (u32, u32, i64, i64) = (8, 6, 0, 4351);
/// `scale_index(v + rvs, 7, 32)` (step 3).
pub const SCALE_INDEX: (u32, i64) = (7, 32);
/// The latent's side, and the hyper-latent's.
const SIDE: u32 = 16;
const Z_SIDE: u32 = 4;

/// A plane of the picture: its channel count, its rows' place in the hyper-latent
/// stream, its scale graph and its tables (§13.6.1, *Each plane*).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Plane {
    Luma,
    Chroma,
}

impl Plane {
    /// `C`: 160 luma channels, 96 chroma.
    pub fn channels(self) -> usize {
        match self {
            Plane::Luma => 160,
            Plane::Chroma => 96,
        }
    }
    /// `skip`: the rows of the hyper-latent stream decoded and discarded before this
    /// plane's own (the luma rows, for chroma).
    pub fn skip(self) -> usize {
        match self {
            Plane::Luma => 0,
            Plane::Chroma => 160,
        }
    }
}

/// A bundle's scale path, lowered and loaded: the two scale graphs, the latent tables,
/// and the tables of `norm.json` the decode reads, after their load checks.
pub struct ScalePath {
    hs: [Loaded; 2],
    tables: Vec<[u32; 256]>,
    bounds: Vec<i64>,
    /// The hyper-latent rows each plane's decode runs, in order: luma's 160, and for
    /// chroma luma's 160 (decoded and discarded) then its own 96.
    rows: [Vec<Vec<i64>>; 2],
    gains: [Vec<i64>; 2],
}

impl ScalePath {
    /// The plane's scale graph, as loaded.
    pub fn graph(&self, p: Plane) -> &Loaded {
        &self.hs[p as usize]
    }
    /// The plane's gains, `gain_y` or `gain_uv`.
    pub fn gains(&self, p: Plane) -> &[i64] {
        &self.gains[p as usize]
    }
    /// What the scale path holds on the heap: the graphs, the tables and the rows.
    pub fn heap_bytes(&self) -> u64 {
        let rows: u64 = self.rows.iter().flatten().map(|r| 24 + 8 * r.capacity() as u64).sum::<u64>()
            + self.rows.iter().map(|r| 24 * r.capacity() as u64).sum::<u64>();
        self.hs.iter().map(|g| g.graph().heap_bytes()).sum::<u64>()
            + 1024 * self.tables.capacity() as u64
            + 8 * self.bounds.capacity() as u64
            + rows
            + self.gains.iter().map(|g| 8 * g.capacity() as u64).sum::<u64>()
    }
}

fn refuse() -> Refusal {
    Refusal::new(Code::ModelMismatch, Site::Lowering)
}

/// A scale network this crate parsed ([`IntNet`]) in the form `wai-ops` lowers it from
/// ([`ScaleNet`]): the members of Table 13.6-A, each as the file gives it. A member the
/// network lacks is refused (`model_mismatch`), as the lowering of the file's bytes
/// refuses it.
pub fn scale_net(net: &IntNet) -> Result<ScaleNet, Refusal> {
    let mut s = ScaleNet::default();
    for (b, names) in SCALE_MEMBERS.iter().enumerate() {
        for (j, name) in names.iter().enumerate() {
            let (v, dims) = net.member(name).ok_or_else(refuse)?;
            s.blocks[b][j] = Member { dims: dims.iter().map(|&d| d as u64).collect(), v: v.to_vec() };
        }
    }
    Ok(s)
}

/// The members of `norm.json` §13.6.1 reads, from the tables this crate parsed
/// ([`NormTables`]).
pub fn norm(n: &NormTables) -> Result<Norm, Refusal> {
    Ok(Norm {
        pmf: n.pmf.clone(),
        bounds: n.bounds.clone(),
        per_channel: i64::try_from(n.per_channel).map_err(|_| refuse())?,
        cdfs_y: n.cdfs_y.clone(),
        cdfs_uv: n.cdfs_uv.clone(),
        gain_y: n.gain_y.clone(),
        gain_uv: n.gain_uv.clone(),
    })
}

/// A bundle's scale path in the form `wai-ops` lowers it from: the parsed tables and the
/// two parsed scale networks.
pub fn net(n: &NormTables, hs_y: &IntNet, hs_uv: &IntNet) -> Result<JpegaiNet, Refusal> {
    Ok(JpegaiNet { norm: norm(n)?, hs_y: scale_net(hs_y)?, hs_uv: scale_net(hs_uv)? })
}

/// Lower and load a bundle's scale path from the bundle as this crate parsed it.
pub fn load(n: &NormTables, hs_y: &IntNet, hs_uv: &IntNet) -> Result<ScalePath, Refusal> {
    load_net(net(n, hs_y, hs_uv)?)
}

/// Lower and load a bundle's scale path from its members' bytes: `norm.json`,
/// `weights/hs_y.json` and `weights/hs_uv.json`.
pub fn load_bytes(norm: &[u8], hs_y: &[u8], hs_uv: &[u8]) -> Result<ScalePath, Refusal> {
    load_net(wai_ops::lower::parse_jpegai(norm, hs_y, hs_uv)?)
}

/// The scale path of a parsed bundle: the load checks and the graphs under [`PROFILE`]
/// (`wai_ops::lower::jpegai_net_loaded`), then the latent tables E8 builds from `pmf`,
/// which the load checks admit.
pub fn load_net(net: JpegaiNet) -> Result<ScalePath, Refusal> {
    let n = net.norm.clone();
    let g = wai_ops::lower::jpegai_net_loaded(net, PROFILE)?;
    let tables = mans::build_latent_tables(&n.pmf).ok_or_else(refuse)?;
    let rows = [fac_rows(&n.cdfs_y, &n.cdfs_uv, Plane::Luma), fac_rows(&n.cdfs_y, &n.cdfs_uv, Plane::Chroma)];
    Ok(ScalePath { hs: [g.hs_y, g.hs_uv], tables, bounds: n.bounds, rows, gains: [n.gain_y, n.gain_uv] })
}

/// One plane's integers (§13.6.1 steps 1–4), each `[C, …]` channel-major: the
/// hyper-latent `z` (`4 × 4` per channel), then at every latent position (`16 × 16`)
/// the scale `v`, the mask, `likely`, the RVS addition, `σ_idx` and the latent `y_q`
/// (0 where not coded).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PlaneInts {
    pub z: Vec<i64>,
    pub v: Vec<i64>,
    pub mask: Vec<bool>,
    pub likely: Vec<i64>,
    pub rvs: Vec<i64>,
    pub sigma: Vec<i64>,
    pub y_q: Vec<i64>,
}

/// What one plane's decode reads besides the codestream: the plane's scale graph,
/// loaded; the hyper-latent rows its `fac` decode runs, `skip` of them before its own
/// (decoded and discarded); and the latent tables and bounds of its `sgm` decode. A
/// [`ScalePath`] gives each plane's ([`ScalePath::setup`]).
#[derive(Clone, Copy)]
pub struct PlaneSetup<'a> {
    pub graph: &'a Loaded,
    pub rows: &'a [Vec<i64>],
    pub skip: usize,
    pub tables: &'a [[u32; 256]],
    pub bounds: &'a [i64],
}

impl ScalePath {
    /// The plane's [`PlaneSetup`].
    pub fn setup(&self, p: Plane) -> PlaneSetup<'_> {
        let k = p as usize;
        PlaneSetup { graph: &self.hs[k], rows: &self.rows[k], skip: p.skip(), tables: &self.tables, bounds: &self.bounds }
    }
}

/// One plane of a picture through §13.6.1: `soz` is the SOZ segment's payload, `latent`
/// the plane's latent segment's (SORP or SORS), and `grfs` the plane's per-channel GRFS
/// flags (one per channel; a flag absent is clear).
pub fn plane(sp: &ScalePath, p: Plane, soz: &[u8], latent: &[u8], grfs: &[bool]) -> Result<PlaneInts, Refusal> {
    plane_with(&sp.setup(p), soz, latent, grfs)
}

/// [`plane`] for the plane `s` describes. Its channel count is the scale graph's.
pub fn plane_with(s: &PlaneSetup<'_>, soz: &[u8], latent: &[u8], grfs: &[bool]) -> Result<PlaneInts, Refusal> {
    let skip = s.skip;
    // 1. The hyper-latent: a fresh `fac` decode over SOZ, the rows before this plane's
    //    (luma's, for chroma) decoded first and discarded.
    let lens = vec![(Z_SIDE * Z_SIDE) as usize; s.rows.len()];
    let syms = mans::decode_fac(soz, s.rows, &lens).map_err(|r| r.refusal())?;
    let z: Vec<i64> = syms.get(skip..).ok_or_else(refuse)?.iter().flatten().map(|&x| x - Z_OFFSET).collect();
    let c = z.len() / (Z_SIDE * Z_SIDE) as usize;
    // 2. The scale graph.
    let v = crate::int_ops::run1(s.graph, Tensor::new(c as u32, Z_SIDE, Z_SIDE, z.clone()), &[])?;
    if (v.c as usize, v.h, v.w) != (c, SIDE, SIDE) {
        return Err(refuse());
    }
    // 3. The stage operators.
    let (b, sh, lo, hi) = LIKELY;
    let likely = stage::block_stat(&v, b, sh, lo, hi).ok_or_else(refuse)?.data;
    let v = v.data;
    let plane = (SIDE * SIDE) as usize;
    let table: Vec<Vec<i64>> = RVS_ADDITIVE.iter().map(|r| r.to_vec()).collect();
    let mut rvs = Vec::with_capacity(v.len());
    // Each channel's mask and `σ_idx` are formed straight into its span of step 4.
    let mut spans: Vec<mans::Span> = Vec::with_capacity(c);
    for ch in 0..c {
        let (mut mask, mut sigma) = (Vec::with_capacity(plane), Vec::with_capacity(plane));
        let flag = grfs.get(ch).copied().unwrap_or(false);
        for i in ch * plane..(ch + 1) * plane {
            let r = stage::band_lookup(likely[i], &RVS_BAND_EDGES, &table, if flag { 0 } else { 1 }).ok_or_else(refuse)?;
            let sum = v[i].checked_add(r).ok_or_else(refuse)?;
            sigma.push(stage::scale_index(sum, SCALE_INDEX.0, SCALE_INDEX.1).ok_or_else(refuse)? as usize);
            mask.push(stage::coded(v[i], CODED_ABOVE));
            rvs.push(r);
        }
        spans.push(mans::Span { mask, sigma });
    }
    // 4. The latent: a fresh `sgm` decode over the plane's segment, one span per channel.
    let y_q = mans::decode_sgm_tables(latent, s.tables, s.bounds, &spans).map_err(|r| r.refusal())?.concat();
    let (mut mask, mut sigma) = (Vec::with_capacity(v.len()), Vec::with_capacity(v.len()));
    for sp in spans {
        mask.extend(sp.mask);
        sigma.extend(sp.sigma.into_iter().map(|x| x as i64));
    }
    Ok(PlaneInts { z, v, mask, likely, rvs, sigma, y_q })
}

/// A plane's scale network in the form the float decode path's bundle gives it
/// (`hyper_scale_decoder_int_weights.json`, one object per plane): three `blocks`, each
/// with `out_ch`, `in_ch`, `k`, `weight`, `bias` and `shift`, and the network's
/// `d2s_block` and `clip_max`. It is Table 13.6-A's network written another way, and is
/// converted to it ([`ScaleNet`]): the weights as `[out_ch, in_ch, k, k]`, the bias as
/// `[out_ch]` and the shifts as `[1, out_ch, 1, 1]`. A network whose depth-to-space block
/// is not 4, whose clip is not 3967, or whose fields are not integers of their kind, is
/// not that network, and is refused (`model_mismatch`).
pub fn scale_net_of_blocks(v: &serde_json::Value) -> Result<ScaleNet, Refusal> {
    let int = |x: &serde_json::Value| x.as_i64().ok_or_else(refuse);
    let ints = |x: &serde_json::Value| -> Result<Vec<i64>, Refusal> { x.as_array().ok_or_else(refuse)?.iter().map(int).collect() };
    let dim = |x: &serde_json::Value| x.as_u64().ok_or_else(refuse);
    if v.get("d2s_block").and_then(serde_json::Value::as_u64) != Some(4) || v.get("clip_max").and_then(serde_json::Value::as_i64) != Some(3967) {
        return Err(refuse());
    }
    let blocks = v.get("blocks").and_then(serde_json::Value::as_array).filter(|b| b.len() == 3).ok_or_else(refuse)?;
    let mut s = ScaleNet::default();
    for (i, b) in blocks.iter().enumerate() {
        let (out, inc, k) = (dim(&b["out_ch"])?, dim(&b["in_ch"])?, dim(&b["k"])?);
        s.blocks[i] = [
            Member { dims: vec![out, inc, k, k], v: ints(&b["weight"])? },
            Member { dims: vec![out], v: ints(&b["bias"])? },
            Member { dims: vec![1, out, 1, 1], v: ints(&b["shift"])? },
        ];
    }
    Ok(s)
}

/// One plane's scale graph from its network alone, loaded under [`PROFILE`]: `cz`
/// hyper-latent channels and the plane's gains (`wai_ops::lower::jpegai::scale`). The
/// calibration tests run one plane's network; a decode loads both through [`load`].
pub fn load_scale(net: ScaleNet, cz: usize, gains: &[i64]) -> Result<Loaded, Refusal> {
    let g = wai_ops::lower::jpegai::scale(net, cz as u64, gains)?;
    wai_ops::graph::load(&g.to_bytes(), PROFILE)
}

/// One `sgm` span of E8 (§5E) over `stream`, with the latent tables built from `pmf`,
/// unsigned `bounds`, and each position's distribution and mask: the reproduction vector
/// the browser crate decodes. Every symbol is a 16-bit value; an unmasked position is 0.
pub fn latent_span(stream: &[u8], pmf: &[Vec<i64>], bounds: &[u8], sigma: &[u8], mask: &[bool]) -> Result<Vec<i16>, Refusal> {
    let bounds: Vec<i64> = bounds.iter().map(|&b| b as i64).collect();
    let span = mans::Span { mask: mask.to_vec(), sigma: sigma.iter().map(|&s| s as usize).collect() };
    let out = mans::decode_sgm(stream, pmf, &bounds, &[span]).map_err(|r| r.refusal())?;
    Ok(out[0].iter().map(|&v| v as i16).collect())
}

/// The rows a plane's `fac` decode runs, from `norm.json`'s `cdfs_y` and `cdfs_uv` (each
/// 63 a row): luma's 160 rows, or for chroma luma's 160 then its own 96.
pub fn fac_rows(cdfs_y: &[i64], cdfs_uv: &[i64], p: Plane) -> Vec<Vec<i64>> {
    let mut rows: Vec<Vec<i64>> = cdfs_y.chunks(63).map(<[i64]>::to_vec).collect();
    if p == Plane::Chroma {
        rows.extend(cdfs_uv.chunks(63).map(<[i64]>::to_vec));
    }
    rows
}

// ------------------------------------------------------------ §13.6.2's networks

/// The profile the other integer networks are loaded under (§13.6.2, §11's
/// `role_profiles`): `acc64`. Their layer normalisations, attention and Q16 activations
/// form values above `i32` whatever their weights (§7.3), so no bundle's networks meet
/// `acc32`.
pub(crate) const NETS_PROFILE: Profile = Profile::Acc64;

/// A bundle's other integer networks, lowered and loaded (§13.6.2): `rec_y` and `rec_uv`
/// (the hyper decoder, and for luma the four context stages), and `syn_y` and `syn_uv`
/// (the synthesis).
pub struct Networks {
    rec: [Loaded; 2],
    syn: [Loaded; 2],
}

impl Networks {
    /// The graph of `role` (`rec_y`, `rec_uv`, `syn_y` or `syn_uv`), as loaded.
    pub fn graph(&self, role: &str) -> Option<&Loaded> {
        match role {
            "rec_y" => Some(&self.rec[0]),
            "rec_uv" => Some(&self.rec[1]),
            "syn_y" => Some(&self.syn[0]),
            "syn_uv" => Some(&self.syn[1]),
            _ => None,
        }
    }
    /// What the networks hold on the heap: the four graphs.
    pub fn heap_bytes(&self) -> u64 {
        self.rec.iter().chain(&self.syn).map(|g| g.graph().heap_bytes()).sum()
    }
}

/// The networks this crate parsed ([`crate::jpegai_synthesis::Weights`]) in the form
/// `wai-ops` lowers them from: the two hyper decoders, the four context stages and the
/// two synthesis networks.
pub fn nets_of(w: &crate::jpegai_decode::DecoderWeights) -> Result<wai_ops::lower::jpegai_nets::JpegaiNets, Refusal> {
    let m = &w.mcm_y;
    Ok(wai_ops::lower::jpegai_nets::JpegaiNets {
        hyp: [w.luma.hyper_decoder.ptq()?, w.chroma.hyper_decoder.ptq()?],
        mcm: [m[0].ptq()?, m[1].ptq()?, m[2].ptq()?, m[3].ptq()?],
        synth: [w.luma.synthesis.ptq()?, w.chroma.synthesis.ptq()?],
    })
}

/// Lower and load the networks of a parsed bundle under [`NETS_PROFILE`].
pub fn load_networks(w: &crate::jpegai_decode::DecoderWeights) -> Result<Networks, Refusal> {
    load_networks_net(nets_of(w)?)
}

/// Lower and load the networks from their files' bytes, in
/// `wai_ops::lower::jpegai_nets::NET_MEMBERS`' order.
pub fn load_networks_bytes(files: [&[u8]; 8]) -> Result<Networks, Refusal> {
    load_networks_net(wai_ops::lower::jpegai_nets::parse_nets(files)?)
}

/// The networks of a parsed bundle, loaded under [`NETS_PROFILE`].
pub fn load_networks_net(n: wai_ops::lower::jpegai_nets::JpegaiNets) -> Result<Networks, Refusal> {
    let [ry, ruv, sy, suv] = wai_ops::lower::jpegai_nets::nets_loaded(n, NETS_PROFILE)?;
    Ok(Networks { rec: [ry, ruv], syn: [sy, suv] })
}

fn run_graph(g: &Loaded, inputs: Vec<Tensor>) -> Result<Tensor, Refusal> {
    let mut out = crate::int_ops::run_graph(g, inputs, &[])?;
    Ok(out.swap_remove(0))
}

/// One plane's reconstructed latent `ŷ [C, 16, 16]` (§13.6.2, `rec_y` or `rec_uv`): the
/// hyper decoder on `z [C, 4, 4]` and, for luma, the four context stages, with the
/// dequantised residual `resi` (Q16, `[C, 16, 16]`). A residual value outside `i32` is
/// refused (`field_range`), and so is any activation outside `i32`.
pub fn reconstruct(n: &Networks, p: Plane, z: Vec<i64>, resi: Vec<i64>) -> Result<Vec<i64>, Refusal> {
    let c = p.channels() as u32;
    let out = run_graph(&n.rec[p as usize], vec![Tensor::new(c, Z_SIDE, Z_SIDE, z), Tensor::new(c, SIDE, SIDE, resi)])?;
    if (out.c, out.h, out.w) != (c, SIDE, SIDE) {
        return Err(refuse());
    }
    Ok(out.data)
}

/// The luma plane `[256, 256]` out of the synthesis, denormalised, in Q16 (§13.6.2,
/// `syn_y`), from the luma latent `[160, 16, 16]` in Q16.
pub fn synthesize_y(n: &Networks, latent_y: Vec<i64>) -> Result<Vec<i64>, Refusal> {
    let out = run_graph(&n.syn[0], vec![Tensor::new(160, SIDE, SIDE, latent_y)])?;
    if (out.c, out.h, out.w) != (1, 256, 256) {
        return Err(refuse());
    }
    Ok(out.data)
}

/// The chroma planes U then V, each `[256, 256]`, out of the synthesis, denormalised, in
/// Q16 (§13.6.2, `syn_uv`), from the luma latent `[160, 16, 16]` and the chroma latent
/// `[96, 16, 16]` in Q16.
pub fn synthesize_uv(n: &Networks, latent_y: Vec<i64>, latent_uv: Vec<i64>) -> Result<Vec<i64>, Refusal> {
    let out = run_graph(&n.syn[1], vec![Tensor::new(160, SIDE, SIDE, latent_y), Tensor::new(96, SIDE, SIDE, latent_uv)])?;
    if (out.c, out.h, out.w) != (2, 256, 256) {
        return Err(refuse());
    }
    Ok(out.data)
}

/// The front half as the exact decoder computed it before it decoded through the
/// operator set (test-only, `legacy_int`): the crate's mANS port, its hyper-scale network
/// and the pipeline's derivations. It panics where its port panics, on a stream that runs
/// out or a table out of range, and it wraps what it reads into narrower types.
#[cfg(any(test, feature = "legacy_int"))]
pub mod legacy {
    use super::{Plane, PlaneInts};
    use crate::jpegai_decode::NormTables;
    use crate::jpegai_mans::{build_decode_transitions, cdf_from_pmf, MansDecoder};
    use crate::jpegai_pipeline::{rvs_additive, rvs_likely, sigma_index, skip_mask};
    use crate::jpegai_synthesis::{decode_hyper_scale, IntNet};

    /// One plane, as `jpegai_decode` decoded it: the same integers [`super::plane`]
    /// gives, wherever neither panics nor refuses.
    pub fn plane(n: &NormTables, hs: &IntNet, gain: &[i64], p: Plane, soz: &[u8], latent: &[u8], grfs: &[bool]) -> PlaneInts {
        const H: usize = 16;
        const W: usize = 16;
        let channels = p.channels();
        let u8s = |v: &[i64]| -> Vec<u8> { v.iter().map(|&x| x as u8).collect() };
        let mut dz = MansDecoder::new(soz, &[soz.len() as u32]);
        if p.skip() > 0 {
            let mut skip = vec![0u8; p.skip() * n.per_channel];
            dz.decode_factorize(&u8s(&n.cdfs_y), &mut skip, p.skip(), n.per_channel);
        }
        let cdfs = match p {
            Plane::Luma => &n.cdfs_y,
            Plane::Chroma => &n.cdfs_uv,
        };
        let mut syms = vec![0u8; channels * n.per_channel];
        dz.decode_factorize(&u8s(cdfs), &mut syms, channels, n.per_channel);
        let z: Vec<i64> = syms.iter().map(|&s| s as i64 - super::Z_OFFSET).collect();
        let (i_s, ..) = decode_hyper_scale(&z, channels, 4, 4, hs);
        let likely = rvs_likely(&i_s, gain, channels, H, W);
        let mut flags = grfs.to_vec();
        flags.resize(channels, false);
        let rvs = rvs_additive(&i_s, gain, &flags, channels, H, W);
        let sigma = sigma_index(&i_s, gain, &rvs, channels);
        let mask = skip_mask(&i_s, gain, channels);
        let flat: Vec<u32> = build_decode_transitions(&n.pmf, &cdf_from_pmf(&n.pmf)).iter().flatten().copied().collect();
        let mut dec = MansDecoder::new(latent, &[latent.len() as u32]);
        dec.set_sgm_transitions(&flat, &u8s(&n.bounds));
        let plane = H * W;
        let mut y_q = vec![0i16; channels * plane];
        for c in 0..channels {
            let mut row = vec![0i16; plane];
            dec.decode_sgm(&sigma[c * plane..(c + 1) * plane], &mut row, &mask[c * plane..(c + 1) * plane]);
            y_q[c * plane..(c + 1) * plane].copy_from_slice(&row);
        }
        let v = i_s.iter().enumerate().map(|(i, &x)| x + gain[i / plane]).collect();
        PlaneInts {
            z,
            v,
            mask,
            likely,
            rvs,
            sigma: sigma.iter().map(|&s| s as i64).collect(),
            y_q: y_q.iter().map(|&q| q as i64).collect(),
        }
    }
}

/// The scale graphs on the host's fast kernels (`crate::int_ops::Host`): a test here and
/// not in `int_ops`, which the byte-equality crate includes without this module.
#[cfg(test)]
mod fast_kernels {
    use super::*;
    use crate::int_kernels::testkit::{fast_conv_runs, Rng, FAST};
    use crate::int_kernels::{with_policy, Policy};

    /// The `wai.image.jpegai` scale path through its lowering `jpegai-int/1` (the committed
    /// bundle's two scale graphs, loaded under `acc32` as `jpegai_ops` loads them): under
    /// every fast policy each graph's three convolutions run on the fast kernel, and every
    /// value is the interpreter's, on hyper-latents over the whole of `z`'s range
    /// `[−31, 32]`, its ends among them.
    #[test]
    fn the_jpegai_scale_path_runs_its_convolutions_on_the_fast_kernel() {
        let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../ops-conformance/sets/jpegai");
        let rd = |n: &str| std::fs::read(dir.join(n)).unwrap();
        let sp = load_bytes(&rd("norm.json"), &rd("weights/hs_y.json"), &rd("weights/hs_uv.json")).unwrap();
        let mut r = Rng(0x0D2C_0007);
        for plane in [Plane::Luma, Plane::Chroma] {
            let c = plane.channels();
            for k in 0..6 {
                let z: Vec<i64> = (0..c * 16).map(|i| match k {
                    0 => -31,
                    1 => 32,
                    _ if i % 7 == 0 => [-31, 32][r.below(2)],
                    _ => r.below(64) as i64 - 31,
                }).collect();
                let x = || Tensor::new(c as u32, 4, 4, z.clone());
                let (want, none) = fast_conv_runs(|| with_policy(Policy::reference(), || crate::int_ops::run1(sp.graph(plane), x(), &[]).unwrap()));
                assert_eq!(none, 0, "the reference policy runs no fast kernel");
                for &p in FAST.iter() {
                    let (got, runs) = fast_conv_runs(|| with_policy(p, || crate::int_ops::run1(sp.graph(plane), x(), &[]).unwrap()));
                    assert!(got == want, "{plane:?} {p:?}: the scale differs from the interpreter alone");
                    assert_eq!(runs, 3, "{plane:?} {p:?}: the three convolutions run on the fast kernel");
                }
            }
        }
    }

    /// The other integer networks through their lowering (§13.6.2: the committed bundle's
    /// four graphs, loaded under `acc64` as `jpegai_ops` loads them): under every fast
    /// policy every convolution and transposed convolution of each graph runs on the fast
    /// kernel, and every value is the interpreter's, on inputs at small sizes over a
    /// range of magnitudes.
    #[test]
    fn the_jpegai_networks_run_their_convolutions_on_the_fast_kernel() {
        let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../ops-conformance/sets/jpegai");
        let files: Vec<Vec<u8>> = wai_ops::lower::jpegai_nets::NET_MEMBERS.iter().map(|n| std::fs::read(dir.join(n)).unwrap()).collect();
        let nets = load_networks_bytes(std::array::from_fn(|i| files[i].as_slice())).unwrap();
        let mut r = Rng(0x0D2C_007B);
        for role in ["rec_y", "rec_uv", "syn_y", "syn_uv"] {
            let mut ran = 0;
            let g = nets.graph(role).unwrap();
            let convs = g.graph().ops.iter().filter(|o| matches!(o.p, wai_ops::graph::Params::Conv { .. } | wai_ops::graph::Params::ConvT { .. })).count() as u64;
            for k in 0..4 {
                let a = 1i64 << (2 + 3 * k);
                let ins: Vec<Tensor> = g.graph().inputs.iter().enumerate().map(|(i, rec)| {
                    let c = g.graph().tensors[rec.tensor as usize].channels;
                    let (h, w) = if role.starts_with("rec") { if i == 0 { (1, 1) } else { (4, 4) } } else { (2, 2) };
                    let n = (c * h * w) as usize;
                    let d = if role.starts_with("rec") && i == 0 { (0..n).map(|_| r.below(64) as i64 - 31).collect() } else { r.vec(n, a, 10) };
                    Tensor::new(c, h, w, d)
                }).collect();
                let run = || wai_ops::run_with(g, ins.clone(), &[], &mut wai_ops::Sequence::new(), &crate::int_ops::Host);
                let (want, none) = fast_conv_runs(|| with_policy(Policy::reference(), run));
                assert_eq!(none, 0, "{role}: the reference policy runs no fast kernel");
                for &p in FAST.iter() {
                    let (got, runs) = fast_conv_runs(|| with_policy(p, run));
                    assert!(got == want, "{role} {p:?}: the outputs differ from the interpreter alone");
                    if want.is_ok() {
                        assert_eq!(runs, convs, "{role} {p:?}: every convolution runs on the fast kernel");
                        ran += 1;
                    }
                }
            }
            assert!(ran > 0, "{role}: some decode completes, so its every operator ran");
        }
    }
}
