//! MLIC++ multi-reference context decoder — the full channel-autoregressive +
//! checkerboard entropy model (Jiang et al., the −13.39%-vs-VVC codec).
//!
//! `elic` ([`super::context`]) gives spatial context within one latent tensor.
//! MLIC++ adds **channel-wise autoregression**: the latent is split into `K`
//! channel slices decoded in order, and *previously decoded slices are context
//! for the next* — its "multiple references" are channel-wise (prior slices),
//! local-spatial (checkerboard), and global-spatial (attention). The attention
//! modules live inside the sink's entropy network; this file implements the
//! decode **orchestration** that drives them:
//!
//!   1. decode `z` → `h_s(z)` = hyperprior params
//!   2. for each channel slice k, in order:
//!        a. params = entropy(concat[ y_decoded_so_far , hyper ])
//!        b. checkerboard PASS A — decode slice-k **anchors** `(y+x) even`
//!        c. params = entropy(…)  (now sees slice-k anchors)
//!        d. checkerboard PASS B — decode slice-k **non-anchors**
//!   3. `g_s(y)` → image
//!
//! Re-running `entropy` after each group is the autoregressive scaffold; the net
//! attends channel-wise to the filled slices and spatially within them. Models
//! are sink-supplied (`hyper_synthesis.onnx`, `entropy.onnx`, `synthesis.onnx`).
//!
//! Wire format (little-endian):
//!   u32 H, u32 W, u32 Lz, u32 Ly, u16 Cz, u16 Cy, u8 Sz_log2, u8 Sy_log2,
//!   u16 q, u16 K
//!   Lz bytes z   |   Ly bytes y  (slice0 anchors, slice0 non-anchors, slice1 …)


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 = 26;

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

/// Decode one channel slice's checkerboard pass into `y`, raster order.
#[allow(clippy::too_many_arguments)]
fn decode_slice_pass(
    dec: &mut RansDecoder,
    y: &mut [f32],
    c0: usize,
    c1: usize,
    sy: usize,
    scales: &[f32],
    means: &[f32],
    g: &GaussCdf,
    anchor: bool,
) {
    let plane = sy * sy;
    for ch in c0..c1 {
        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];
            }
        }
    }
}

/// Build the `entropy` net input: `concat([y, hyper], dim=channels)` →
/// `[1, 3*Cy, Sy, Sy]` (`y` is `Cy` channels, `hyper` is `2*Cy`).
fn entropy_input(y: &[f32], hyper: &[f32], cy: usize, sy: usize) -> Array4<f32> {
    let plane = sy * sy;
    let mut t = Array4::<f32>::zeros((1, 3 * cy, sy, sy));
    let s = t.as_slice_mut().unwrap();
    s[0..cy * plane].copy_from_slice(y);
    s[cy * plane..3 * cy * plane].copy_from_slice(hyper);
    t
}

pub(crate) fn decode(payload: &[u8], params: &ParamSource) -> Result<DecodedImage, DecodeError> {
    if payload.len() < HEADER {
        return Err(DecodeError::InvalidPayload(format!(
            "mlicpp 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;
    let k = u16a(24) as usize;
    if payload.len() != HEADER + lz + ly {
        return Err(DecodeError::InvalidPayload(format!(
            "mlicpp payload {} != {HEADER}+Lz={lz}+Ly={ly}", payload.len())));
    }
    if k == 0 || cy % k != 0 {
        return Err(DecodeError::InvalidPayload(format!("mlicpp Cy={cy} not divisible by K={k}")));
    }
    let slice_w = cy / k;
    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. hyperprior params (the channel-independent base reference).
    let (hyper, hd) = run_onnx(params, "hyper_synthesis.onnx", z, 2 * cy)?;
    if hd[1] != 2 * cy || hd[2] != sy || hd[3] != sy {
        return Err(DecodeError::Ort(format!(
            "hyper_synthesis [1,{},{},{}]; want [1,{},{sy},{sy}]", hd[1], hd[2], hd[3], 2 * cy)));
    }
    let plane = sy * sy;
    let g: GaussCdf = read_json(params, &format!("gauss_cdfs_q{q}"))?;
    let entropy = "entropy.onnx";

    // 3. channel-autoregressive slices × checkerboard, one y rANS stream.
    let mut y = vec![0f32; cy * plane];
    let mut dy = RansDecoder::new(y_bits).map_err(|e| DecodeError::InvalidPayload(e.into()))?;
    for k_i in 0..k {
        let (c0, c1) = (k_i * slice_w, (k_i + 1) * slice_w);
        for anchor in [true, false] {
            let (ent_out, pd) = run_onnx(params, entropy, entropy_input(&y, &hyper, cy, sy), 2 * cy)?;
            if pd[1] != 2 * cy {
                return Err(DecodeError::Ort(format!(
                    "entropy out ch {} != 2*Cy={}", pd[1], 2 * cy)));
            }
            let (scales, means) = (&ent_out[0..cy * plane], &ent_out[cy * plane..2 * cy * plane]);
            decode_slice_pass(&mut dy, &mut y, c0, c1, sy, scales, means, &g, anchor);
        }
    }

    // 4. 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::*;

    #[test]
    fn slices_and_passes_visit_every_element_once() {
        // K channel slices × 2 checkerboard passes must tile (Cy × Sy × Sy)
        // exactly — no element decoded twice, none missed.
        let (cy, sy, k) = (8usize, 4usize, 4usize);
        let slice_w = cy / k;
        let plane = sy * sy;
        let mut hits = vec![0u32; cy * plane];
        for k_i in 0..k {
            let (c0, c1) = (k_i * slice_w, (k_i + 1) * slice_w);
            for anchor in [true, false] {
                for ch in c0..c1 {
                    for yy in 0..sy {
                        for xx in 0..sy {
                            if is_anchor(yy, xx) == anchor {
                                hits[ch * plane + yy * sy + xx] += 1;
                            }
                        }
                    }
                }
            }
        }
        assert!(hits.iter().all(|&n| n == 1), "every latent element decoded exactly once");
    }
}
