//! GLC — Generative Latent Coding (`wai.neural.glc`), ultra-low-bitrate image
//! compression in a VQ-VAE latent space.
//!
//! GLC is NOT the Gaussian-conditional family: its latent is a grid of
//! **categorical** codebook indices, and its entropy model predicts a
//! per-position categorical distribution (softmax over the codebook) from a
//! hyperprior + spatial context. So this decoder is categorical, not Gaussian:
//!
//!   z (factorized)  →  hyper_entropy → base logits[K] per position
//!   PASS A (anchors)     : decode index from softmax(base)
//!   context_entropy(anchors) → ctx logits
//!   PASS B (non-anchors) : decode index from softmax(base + ctx)
//!   synthesis(index grid) → image      (codebook lookup + generative decoder)
//!
//! The categorical CDF is built per position with [`pmf_to_quantized_cdf`] —
//! the exact CompressAI quantizer (round → renorm → cumsum → steal-from-richest
//! so no symbol has zero frequency) — and decoded with
//! [`RansDecoder::decode_categorical`]. As with every learned entropy model,
//! encoder and decoder MUST compute the logits and quantize identically (the
//! determinism contract); GLC's float softmax makes that sharper than the
//! hyperprior's scale buckets.
//!
//! Wire format (little-endian):
//!   <IIIIHHBBH>  H, W, Lz, Ly, Cz, K, Sz_log2, Sy_log2, q
//!   Lz bytes z-bitstream | Ly bytes index-bitstream
//! Sink ONNX: hyper_entropy.onnx (z→logits[1,K,Sy,Sy]),
//!            context_entropy.onnx (anchor-index grid [1,1,Sy,Sy]→logits),
//!            synthesis.onnx (index grid [1,1,Sy,Sy]→image), + hyper_cdfs json.


use ndarray::Array4;

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

const HEADER: usize = 24;
const PRECISION: u32 = 16;

/// CompressAI's `pmf_to_quantized_cdf`: a float PMF → an integer CDF summing to
/// `2^precision`, with every symbol guaranteed ≥ 1 frequency. Returns
/// `cdf[0..=K]` (length `K+1`), `cdf[0]=0`, `cdf[K]=2^precision`. Bit-exact with
/// the C++ extension — unit-tested against it.
pub(crate) fn pmf_to_quantized_cdf(pmf: &[f32], precision: u32) -> Vec<u32> {
    let k = pmf.len();
    let mut cdf = vec![0u32; k + 1];
    // 1. per-symbol frequency = round(p * 2^precision), stored at cdf[1..].
    for (i, &p) in pmf.iter().enumerate() {
        cdf[i + 1] = (p * (1u64 << precision) as f32).round() as u32;
    }
    // 2. renorm the frequencies to sum to 2^precision.
    let total: u64 = cdf.iter().map(|&v| v as u64).sum();
    let total = total.max(1);
    for v in &mut cdf {
        *v = (((1u64 << precision) * (*v as u64)) / total) as u32;
    }
    // 3. cumulative, then pin the tail exactly.
    for i in 1..cdf.len() {
        cdf[i] += cdf[i - 1];
    }
    *cdf.last_mut().unwrap() = 1u32 << precision;
    // 4. fix any zero-frequency symbol by stealing from the richest.
    for i in 0..cdf.len() - 1 {
        if cdf[i] == cdf[i + 1] {
            let (mut best_freq, mut best) = (u32::MAX, usize::MAX);
            for j in 0..cdf.len() - 1 {
                let freq = cdf[j + 1] - cdf[j];
                if freq > 1 && freq < best_freq {
                    best_freq = freq;
                    best = j;
                }
            }
            if best == usize::MAX {
                continue;
            }
            if best < i {
                for v in cdf.iter_mut().take(i + 1).skip(best + 1) {
                    *v -= 1;
                }
            } else {
                for v in cdf.iter_mut().take(best + 1).skip(i + 1) {
                    *v += 1;
                }
            }
        }
    }
    cdf
}

fn softmax(logits: &[f32]) -> Vec<f32> {
    let m = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    let mut e: Vec<f32> = logits.iter().map(|&l| (l - m).exp()).collect();
    let s: f32 = e.iter().sum::<f32>().max(f32::MIN_POSITIVE);
    for v in &mut e {
        *v /= s;
    }
    e
}

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

/// Decode one position's categorical index from per-class `logits` (length K).
fn decode_index(dec: &mut RansDecoder, logits: &[f32]) -> u32 {
    let cdf = pmf_to_quantized_cdf(&softmax(logits), PRECISION);
    dec.decode_categorical(&cdf)
}

pub(crate) fn decode(payload: &[u8], params: &ParamSource) -> Result<DecodedImage, DecodeError> {
    if payload.len() < HEADER {
        return Err(DecodeError::InvalidPayload(format!(
            "glc 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 = u16a(16) as usize;
    let k = 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!(
            "glc payload {} != {HEADER}+Lz={lz}+Ly={ly}", payload.len())));
    }
    let z_bits = &payload[HEADER..HEADER + lz];
    let y_bits = &payload[HEADER + lz..HEADER + lz + ly];
    let plane = sy * sy;

    // 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;
            }
        }
    }

    // base logits[K] per position from the hyperprior.
    let (base, bd) = run_onnx(params, "hyper_entropy.onnx", z, k)?;
    if bd[1] != k || bd[2] != sy || bd[3] != sy {
        return Err(DecodeError::Ort(format!(
            "hyper_entropy [1,{},{},{}]; want [1,{k},{sy},{sy}]", bd[1], bd[2], bd[3])));
    }

    let mut idx = vec![0u32; plane];
    let mut dy = RansDecoder::new(y_bits).map_err(|e| DecodeError::InvalidPayload(e.into()))?;
    let logits_at = |t: &[f32], pos: usize| -> Vec<f32> {
        (0..k).map(|c| t[c * plane + pos]).collect()
    };

    // PASS A — anchors, base logits.
    for yy in 0..sy {
        for xx in 0..sy {
            if is_anchor(yy, xx) {
                let pos = yy * sy + xx;
                idx[pos] = decode_index(&mut dy, &logits_at(&base, pos));
            }
        }
    }

    // context from the anchor index grid.
    let mut anchor_grid = Array4::<f32>::zeros((1, 1, sy, sy));
    {
        let g = anchor_grid.as_slice_mut().unwrap();
        for yy in 0..sy {
            for xx in 0..sy {
                if is_anchor(yy, xx) {
                    g[yy * sy + xx] = idx[yy * sy + xx] as f32;
                }
            }
        }
    }
    let (ctx, cd) = run_onnx(params, "context_entropy.onnx", anchor_grid, k)?;
    if cd[1] != k || cd[2] != sy || cd[3] != sy {
        return Err(DecodeError::Ort(format!(
            "context_entropy [1,{},{},{}]; want [1,{k},{sy},{sy}]", cd[1], cd[2], cd[3])));
    }

    // PASS B — non-anchors, base + context logits.
    for yy in 0..sy {
        for xx in 0..sy {
            if !is_anchor(yy, xx) {
                let pos = yy * sy + xx;
                let logits: Vec<f32> =
                    (0..k).map(|c| base[c * plane + pos] + ctx[c * plane + pos]).collect();
                idx[pos] = decode_index(&mut dy, &logits);
            }
        }
    }

    // synthesis: index grid → image (codebook + generative decoder in the sink).
    let mut grid = Array4::<f32>::zeros((1, 1, sy, sy));
    grid.as_slice_mut().unwrap().iter_mut().zip(&idx).for_each(|(g, &v)| *g = v as f32);
    let (img, id) = run_onnx(params, "synthesis.onnx", grid, 3)?;
    if id[2] != h || id[3] != w {
        return Err(DecodeError::Ort(format!(
            "glc 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 p in 0..hw {
        for ch in 0..3 {
            rgb[p * 3 + ch] = (img[ch * hw + p].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 quantized_cdf_is_bit_exact_with_compressai() {
        // Reference vectors from CompressAI's C++ `pmf_to_quantized_cdf`.
        let cases: &[(&[f32], &[u32])] = &[
            (&[0.25, 0.25, 0.25, 0.25], &[0, 16384, 32768, 49152, 65536]),
            (&[0.97, 0.01, 0.01, 0.01], &[0, 63570, 64225, 64880, 65536]),
            (&[0.5, 0.5], &[0, 32768, 65536]),
            (&[0.001, 0.001, 0.001, 0.997], &[0, 65, 130, 195, 65536]),
        ];
        for (pmf, expect) in cases {
            assert_eq!(&pmf_to_quantized_cdf(pmf, PRECISION), expect, "pmf {pmf:?}");
        }
        // arange(1..=8)/36 — exercises the rescale path.
        let pmf: Vec<f32> = (1..=8).map(|i| i as f32 / 36.0).collect();
        assert_eq!(pmf_to_quantized_cdf(&pmf, PRECISION),
                   vec![0, 1820, 5461, 10922, 18204, 27306, 38229, 50972, 65536]);
    }

    #[test]
    fn quantized_cdf_is_monotone_sums_to_total_no_zero_freq() {
        for pmf in [
            vec![0.25f32, 0.25, 0.25, 0.25],
            vec![0.97, 0.01, 0.01, 0.01],
            vec![0.5, 0.5],
            vec![0.001, 0.001, 0.001, 0.997],
            (0..64).map(|i| (i + 1) as f32).collect::<Vec<_>>(),
        ] {
            let s: f32 = pmf.iter().sum();
            let pmf: Vec<f32> = pmf.iter().map(|p| p / s).collect();
            let cdf = pmf_to_quantized_cdf(&pmf, PRECISION);
            assert_eq!(cdf.len(), pmf.len() + 1);
            assert_eq!(cdf[0], 0);
            assert_eq!(*cdf.last().unwrap(), 1 << PRECISION);
            for i in 0..cdf.len() - 1 {
                assert!(cdf[i + 1] > cdf[i], "zero/negative freq at {i}: {cdf:?}");
            }
        }
    }
}
