//! Checkerboard **context-model** image decoder — the architectural class of
//! JPEG AI and MLIC++, a tier above the plain hyperprior.
//!
//! The hyperprior ([`super::hyperprior`]) gives each latent a (scale, mean) from
//! the side stream `z` alone. A *context model* additionally conditions each
//! latent on already-decoded SPATIAL NEIGHBORS — which is most of the coding
//! gain in modern learned codecs, and exactly what the ISO JPEG AI entropy
//! engine does. We implement the **checkerboard** variant (He et al. 2021 /
//! ELIC): a 2-pass scheme that gets spatial context with only TWO network
//! passes instead of one-per-pixel autoregression.
//!
//!   1. decode `z` (factorized) → `h_s(z)` → base (scale, mean)
//!   2. PASS A — decode the **anchor** latents `(y+x) even` with base params
//!   3. run `context(y_anchors)` → refined (scale, mean) for the rest
//!   4. PASS B — decode the **non-anchor** latents `(y+x) odd` with those
//!   5. `g_s(y_hat)` → image
//!
//! This is a real, WAI-native codec of the JPEG-AI/MLIC++ entropy class — NOT a
//! reimplementation of the ISO JPEG AI bitstream (that needs the standard text +
//! reference weights). Models are sink-supplied, like every neural capability.
//!
//! Wire format (little-endian, == hyperprior):
//!   u32 H, u32 W, u32 Lz, u32 Ly, u16 Cz, u16 Cy, u8 Sz_log2, u8 Sy_log2, u16 q
//!   Lz bytes z bitstream   |   Ly bytes y bitstream (anchors then non-anchors,
//!                                                     channel-major, raster)
//! Sink files: hyper_synthesis.onnx, context.onnx, synthesis.onnx,
//!             hyper_cdfs_q<q>.json, gauss_cdfs_q<q>.json.


use ndarray::Array4;

use super::hyperprior::{build_index, read_json, run_onnx, FactorizedCdf, GaussCdf};
use super::image::DecodedImage;
use super::{params::ParamSource, rans::RansDecoder, DecodeError};

const HEADER: usize = 24;

#[inline]
fn is_anchor(yy: usize, xx: usize) -> bool {
    (yy + xx).is_multiple_of(2)
}

/// Decode the latent positions selected by `anchor` (a checkerboard pass) into
/// `y`, in channel-major raster order, using per-element (scale, mean).
#[allow(clippy::too_many_arguments)]
fn decode_pass(
    dec: &mut RansDecoder,
    y: &mut [f32],
    cy: usize,
    sy: usize,
    scales: &[f32],
    means: &[f32],
    g: &GaussCdf,
    anchor: bool,
) {
    let plane = sy * sy;
    for ch in 0..cy {
        for yy in 0..sy {
            for xx in 0..sy {
                if is_anchor(yy, xx) != anchor {
                    continue;
                }
                let e = ch * plane + yy * sy + xx;
                let b = build_index(scales[e], &g.scale_table);
                let sym = dec.decode(&g.cdf[b], g.cdf_length[b], g.offset[b]);
                y[e] = sym as f32 + means[e];
            }
        }
    }
}

pub(crate) fn decode(payload: &[u8], params: &ParamSource) -> Result<DecodedImage, DecodeError> {
    if payload.len() < HEADER {
        return Err(DecodeError::InvalidPayload(format!(
            "context payload {} B < {HEADER} B header", payload.len())));
    }
    let u32a = |o: usize| u32::from_le_bytes(payload[o..o + 4].try_into().unwrap());
    let u16a = |o: usize| u16::from_le_bytes(payload[o..o + 2].try_into().unwrap());
    let (h, w) = (u32a(0) as usize, u32a(4) as usize);
    let (lz, ly) = (u32a(8) as usize, u32a(12) as usize);
    let (cz, cy) = (u16a(16) as usize, u16a(18) as usize);
    let (sz, sy) = (1usize << payload[20], 1usize << payload[21]);
    let q = u16a(22) as u32;
    if payload.len() != HEADER + lz + ly {
        return Err(DecodeError::InvalidPayload(format!(
            "context payload {} != {HEADER}+Lz={lz}+Ly={ly}", payload.len())));
    }
    let z_bits = &payload[HEADER..HEADER + lz];
    let y_bits = &payload[HEADER + lz..HEADER + lz + ly];

    // 1. z (factorized).
    let hcdf: FactorizedCdf = read_json(params, &format!("hyper_cdfs_q{q}"))?;
    if hcdf.channels != cz {
        return Err(DecodeError::InvalidPayload(format!("z Cz={cz} != cdf {}", hcdf.channels)));
    }
    let mut z = Array4::<f32>::zeros((1, cz, sz, sz));
    {
        let zs = z.as_slice_mut().unwrap();
        let mut dz = RansDecoder::new(z_bits).map_err(|e| DecodeError::InvalidPayload(e.into()))?;
        let mut i = 0;
        for ch in 0..cz {
            for _ in 0..(sz * sz) {
                zs[i] = dz.decode(&hcdf.cdf[ch], hcdf.cdf_length[ch], hcdf.offset[ch]) as f32
                    + hcdf.medians[ch];
                i += 1;
            }
        }
    }

    // 2. base params from the hyperprior.
    let (base, bd) = run_onnx(params, "hyper_synthesis.onnx", z, 2 * cy)?;
    if bd[1] != 2 * cy || bd[2] != sy || bd[3] != sy {
        return Err(DecodeError::Ort(format!(
            "hyper_synthesis [1,{},{},{}]; want [1,{},{sy},{sy}]", bd[1], bd[2], bd[3], 2 * cy)));
    }
    let plane = sy * sy;
    let (base_scales, base_means) = (&base[0..cy * plane], &base[cy * plane..2 * cy * plane]);

    let g: GaussCdf = read_json(params, &format!("gauss_cdfs_q{q}"))?;
    let mut y = vec![0f32; cy * plane];

    // 3. PASS A — anchors, base params, single y stream.
    let mut dy = RansDecoder::new(y_bits).map_err(|e| DecodeError::InvalidPayload(e.into()))?;
    decode_pass(&mut dy, &mut y, cy, sy, base_scales, base_means, &g, true);

    // 4. context network sees the anchor-only latent, predicts the rest.
    let mut y_anchor = Array4::<f32>::zeros((1, cy, sy, sy));
    y_anchor.as_slice_mut().unwrap().copy_from_slice(&y);
    let (ctx, cd) = run_onnx(params, "context.onnx", y_anchor, 2 * cy)?;
    if cd[1] != 2 * cy || cd[2] != sy || cd[3] != sy {
        return Err(DecodeError::Ort(format!(
            "context [1,{},{},{}]; want [1,{},{sy},{sy}]", cd[1], cd[2], cd[3], 2 * cy)));
    }
    let (ctx_scales, ctx_means) = (&ctx[0..cy * plane], &ctx[cy * plane..2 * cy * plane]);

    // 5. PASS B — non-anchors, context params, SAME y stream continues.
    decode_pass(&mut dy, &mut y, cy, sy, ctx_scales, ctx_means, &g, false);

    // 6. synthesis.
    let mut y4 = Array4::<f32>::zeros((1, cy, sy, sy));
    y4.as_slice_mut().unwrap().copy_from_slice(&y);
    let (img, id) = run_onnx(params, "synthesis.onnx", y4, 3)?;
    if id[2] != h || id[3] != w {
        return Err(DecodeError::Ort(format!(
            "synthesis [1,3,{},{}]; want [1,3,{h},{w}]", id[2], id[3])));
    }
    let mut rgb = vec![0u8; h * w * 3];
    let hw = h * w;
    for yy in 0..h {
        for xx in 0..w {
            let pi = yy * w + xx;
            for ch in 0..3 {
                rgb[pi * 3 + ch] = (img[ch * hw + pi].clamp(0.0, 1.0) * 255.0).round() as u8;
            }
        }
    }
    Ok(DecodedImage { rgb, width: w as u32, height: h as u32 })
}

#[cfg(test)]
mod tests {
    use super::is_anchor;

    #[test]
    fn checkerboard_partitions_exactly() {
        // Every position is in exactly one pass; the two passes tile the grid.
        for s in [2usize, 4, 8, 16] {
            let mut anchors = 0;
            let mut non = 0;
            for y in 0..s {
                for x in 0..s {
                    if is_anchor(y, x) {
                        anchors += 1;
                    } else {
                        non += 1;
                    }
                    // anchor-ness is exactly the (y+x) parity — no overlap.
                    assert_eq!(is_anchor(y, x), !((y + x) % 2 == 1));
                }
            }
            assert_eq!(anchors + non, s * s);
            // A square even grid splits 50/50.
            assert_eq!(anchors, s * s / 2);
            // Neighbors differ in pass (the property the context model exploits).
            assert_ne!(is_anchor(0, 0), is_anchor(0, 1));
            assert_ne!(is_anchor(0, 0), is_anchor(1, 0));
        }
    }
}
