//! Mean-scale **hyperprior** learned-image decoder — the conditional-entropy
//! family (`mbt2018-mean`, `cheng2020`, and the base MLIC++ builds on), a real
//! step past bmshj2018's *factorized* prior.
//!
//! bmshj2018 decodes every latent against ONE static per-channel CDF. A
//! hyperprior instead conditions each latent on a side stream: decode the
//! hyper-latent `z`, run `h_s(z)` to get a per-element **(mean, scale)**, and
//! decode the main latent `y` with a **Gaussian-conditional** model — the CDF
//! per element is chosen by its scale, and its mean is added back. Then
//! `g_s(y)` synthesizes the image. This is the CompressAI decode, exactly.
//!
//! Wire format (little-endian):
//!   u32 H, u32 W, u32 Lz, u32 Ly, u16 Cz, u16 Cy, u8 Sz_log2, u8 Sy_log2, u16 q
//!   Lz bytes  z rANS bitstream   (factorized, decoded with `hyper_cdfs_q<q>`)
//!   Ly bytes  y rANS bitstream   (Gaussian-conditional, `gauss_cdfs_q<q>`)
//!
//! CRITICAL — encoder/decoder must compute the entropy params identically.
//! The scale `σ` selects a CDF *bucket* (`build_index`); a single bucket
//! mismatch corrupts the rANS state and cascades. Since this decoder gets `σ`
//! from the ONNX `h_s`, the ENCODER must too — using torch `h_s` instead flips
//! a few boundary buckets (~1e-5 float drift) and the decode diverges. Verified
//! end-to-end against the real CompressAI mbt2018-mean: torch-encode → 1188
//! corrupt latents; ONNX-encode → max|Δ|=1 (only g_s ORT-version rounding).
//! This is the determinism contract in [`crate::container::DeterminismClass`],
//! live: float entropy-param computation is a cross-tool hazard.
//!
//! Sink files (next to the registered model path):
//!   hyper_cdfs_q<q>.json   factorized z CDF table (same shape as bmshj2018's)
//!   gauss_cdfs_q<q>.json   { scale_table, cdf[], cdf_length[], offset[] }
//!   hyper_synthesis.onnx   z -> [1, 2*Cy, Sy, Sy] (scales ‖ means, CompressAI order)
//!   synthesis.onnx         y_hat -> [1, 3, H, W]


use ndarray::Array4;
use ort::value::Value;
use serde::Deserialize;

use super::image::DecodedImage;
use super::{params::ParamSource, rans::RansDecoder, DecodeError};

const HEADER: usize = 24; // 4+4+4+4 + 2+2 + 1+1 + 2

/// Factorized CDF table (the hyper-latent `z` uses this — one CDF per channel).
/// Shared with the context-model decoder.
#[derive(Debug, Clone, Deserialize)]
pub(crate) struct FactorizedCdf {
    pub(crate) channels: usize,
    pub(crate) cdf_length: Vec<usize>,
    pub(crate) offset: Vec<i32>,
    pub(crate) medians: Vec<f32>,
    pub(crate) cdf: Vec<Vec<i32>>,
}

/// Gaussian-conditional CDF table — one CDF per scale bucket, selected per
/// element by [`build_index`]. No medians: the mean is per-element from `h_s`.
#[derive(Debug, Clone, Deserialize)]
pub(crate) struct GaussCdf {
    /// Ascending scale (σ) bucket edges; an element's σ picks a bucket.
    pub(crate) scale_table: Vec<f32>,
    pub(crate) cdf_length: Vec<usize>,
    pub(crate) offset: Vec<i32>,
    pub(crate) cdf: Vec<Vec<i32>>,
}

pub(crate) fn read_json<T: for<'de> Deserialize<'de>>(params: &ParamSource, stem: &str) -> Result<T, DecodeError> {
    params.json(stem)
}

/// CompressAI `GaussianConditional.build_indexes`: bucket a σ to a scale-table
/// index. Index rises with σ; clamped to `[0, len-1]`. Pure logic — unit-tested.
/// Shared with the context-model decoder.
pub(crate) fn build_index(sigma: f32, scale_table: &[f32]) -> usize {
    let n = scale_table.len();
    // lower_bound: never below the smallest tabulated scale.
    let sigma = sigma.max(scale_table[0]);
    let mut idx = n - 1;
    for &s in &scale_table[..n - 1] {
        if sigma <= s {
            idx -= 1;
        }
    }
    idx
}

pub(crate) fn run_onnx(params: &ParamSource, model: &str, input: Array4<f32>, want_c: usize) -> Result<(Vec<f32>, [usize; 4]), DecodeError> {
    let mut sess = params.session_sibling(model)?;
    let in_name = sess.inputs().first().map(|i| i.name().to_string()).unwrap_or_else(|| "input".into());
    let out_name = sess.outputs().first().map(|o| o.name().to_string()).unwrap_or_else(|| "output".into());
    let v = Value::from_array(input).map_err(|e| DecodeError::Ort(format!("tensor: {e}")))?;
    let outputs = sess.run(ort::inputs![in_name => v])
        .map_err(|e| DecodeError::Ort(format!("{}: {e}", model)))?;
    let out = outputs.get(out_name.as_str())
        .ok_or_else(|| DecodeError::Ort(format!("{} missing output", model)))?;
    let view = out.try_extract_array::<f32>().map_err(|e| DecodeError::Ort(format!("extract: {e}")))?;
    let view = view.view();
    let d = view.shape();
    if d.len() != 4 || (want_c != 0 && d[1] != want_c) {
        return Err(DecodeError::Ort(format!(
            "{} output shape {d:?}; want [1, {}, _, _]", model,
            if want_c == 0 { "C".into() } else { want_c.to_string() })));
    }
    Ok((view.as_slice().expect("contiguous").to_vec(), [d[0], d[1], d[2], d[3]]))
}

/// Entropy-decode a mean-scale-hyperprior latent: factorized `z` → `h_s(z)` →
/// per-element (scale, mean) → Gaussian-conditional `y`. Returns `y_hat`
/// `[1, Cy, Sy, Sy]` — the input to a synthesis (still image) OR a temporal
/// decoder (DCVC video). Shared so the video frame loop reuses the exact,
/// CompressAI-validated entropy path.
#[allow(clippy::too_many_arguments)]
pub(crate) fn decode_yhat(
    z_bits: &[u8], y_bits: &[u8], params: &ParamSource,
    q: u32, cz: usize, cy: usize, sz: usize, sy: usize,
) -> Result<Array4<f32>, DecodeError> {
    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;
            }
        }
    }
    let (hyper_out, pd) = run_onnx(params, "hyper_synthesis.onnx", z, 2 * cy)?;
    if pd[1] / 2 != cy || pd[2] != sy || pd[3] != sy {
        return Err(DecodeError::Ort(format!(
            "hyper_synthesis gave [1,{},{},{}]; want means‖scales for [1,{cy},{sy},{sy}]",
            pd[1], pd[2], pd[3])));
    }
    let plane = sy * sy;
    let scales = &hyper_out[0..cy * plane];
    let means = &hyper_out[cy * plane..2 * cy * plane];
    let gcdf: GaussCdf = read_json(params, &format!("gauss_cdfs_q{q}"))?;
    let mut y = Array4::<f32>::zeros((1, cy, sy, sy));
    {
        let ys = y.as_slice_mut().unwrap();
        let mut dy = RansDecoder::new(y_bits).map_err(|e| DecodeError::InvalidPayload(e.into()))?;
        for e in 0..cy * plane {
            let bucket = build_index(scales[e], &gcdf.scale_table);
            let sym = dy.decode(&gcdf.cdf[bucket], gcdf.cdf_length[bucket], gcdf.offset[bucket]);
            ys[e] = sym as f32 + means[e];
        }
    }
    Ok(y)
}

pub(crate) fn decode(payload: &[u8], params: &ParamSource) -> Result<DecodedImage, DecodeError> {
    if payload.len() < HEADER {
        return Err(DecodeError::InvalidPayload(format!(
            "payload too short ({} B); need ≥ {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!(
            "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-3. Factorized z → h_s → Gaussian-conditional y (shared with video).
    let y = decode_yhat(z_bits, y_bits, params, q, cz, cy, sz, sy)?;

    // 4. g_s(y_hat) -> image.
    let (img, id) = run_onnx(params, "synthesis.onnx", y, 3)?;
    if id[2] != h || id[3] != w {
        return Err(DecodeError::Ort(format!(
            "synthesis gave [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::build_index;

    #[test]
    fn scale_bucketing_is_monotonic_and_clamped() {
        // CompressAI default-style ascending scale table.
        let st = vec![0.11, 0.22, 0.44, 0.88, 1.76, 3.5, 7.0];
        // Below the floor → bucket 0.
        assert_eq!(build_index(0.001, &st), 0);
        assert_eq!(build_index(0.11, &st), 0);
        // Above the top → last bucket.
        assert_eq!(build_index(100.0, &st), st.len() - 1);
        // Monotonic non-decreasing in σ.
        let mut last = 0;
        for i in 0..400 {
            let sigma = i as f32 * 0.02;
            let b = build_index(sigma, &st);
            assert!(b >= last, "bucket dropped at σ={sigma}");
            assert!(b < st.len());
            last = b;
        }
        // A mid σ lands in a middle bucket.
        assert!(build_index(0.5, &st) > 0 && build_index(0.5, &st) < st.len() - 1);
    }
}
