//! Integer Gaussian-conditional entropy — the integer-entropy contract's
//! reusable core (SPEC §7, *Replicate-soundness*).
//!
//! A learned conditional codec entropy-decodes each latent against a CDF chosen
//! by that element's scale `σ`. In the float path ([`crate::neural`]'s
//! `hyperprior::build_index`) `σ` is a float from an ONNX hyper-decoder and the
//! bucket is a float comparison, so ~1e-5 of cross-platform float drift in `σ`
//! flips a boundary bucket and desyncs the **entire** stream. That single fact
//! is why the conditional-entropy family is
//! [`crate::codecs::ReplicateClass::NeuralFloat`]: even with an integer
//! arithmetic coder downstream, a float σ source makes the decode reproducible
//! only on the same build.
//!
//! This module is the integer replacement for that bucketing step: `σ` carried
//! as a fixed-point integer, bucketed by integer comparison. It is deterministic
//! **by construction** — given the same integer `σ` it returns the same bucket
//! on every machine, because there is no float operation on the path. The
//! existing integer rANS is the arithmetic coder these indices feed; the integer
//! synthesis ([`crate::int_synth`]) is the other half.
//!
//! ## What graduates a codec to `NeuralIntegerExact`
//!
//! Adopting integer bucketing is necessary but not sufficient. The desync lives
//! in the *source* of `σ`, not only in the comparison. A codec is cross-hardware
//! replicate-sound when **all three** hold:
//!
//! 1. `σ` is produced by an **integer** hyper-decoder, so the σ bytes are
//!    identical on every machine (the reference integer source is JPEG AI's scale
//!    path through the integer operator set, `jpegai_ops`, whose `scale_index`
//!    (integer-operators §5E E5) is the integer-log-domain special case of
//!    [`det_build_index`]);
//! 2. the scale→bucket step uses this module (no float comparison);
//! 3. the synthesis is integer ([`crate::int_synth`]).
//!
//! With a **float** ONNX hyper-decoder, [`quantize_scale`] snaps σ to the grid
//! and removes drift *inside a bucket*, but cannot guarantee the boundary case —
//! so such a codec stays `NeuralFloat`. Pure i64, no ONNX: like `int_synth`,
//! this builds and is checkable for bit-exactness without `ort`, and can compile
//! to WASM for an in-browser deterministic decode.

/// Precompute a CompressAI-style ascending `scale_table` as fixed-point integer
/// bucket edges (Q`frac`: `frac` fractional bits, grid step `2^-frac`). Compute
/// once per model; pass the result to [`det_build_index`].
pub fn fixed_scale_edges(scale_table: &[f32], frac: u32) -> Vec<i64> {
    let s = (1u64 << frac) as f64;
    scale_table.iter().map(|&e| (e as f64 * s).floor() as i64).collect()
}

/// Snap a float `σ` to the fixed-point grid (Q`frac`), deterministic floor.
///
/// For a codec whose σ source is still a float ONNX hyper-decoder: this removes
/// drift smaller than the grid step *within* a bucket, but is **not** a hard
/// guarantee at a bucket edge (a σ within one grid step of an edge can still
/// land differently across platforms). A guaranteed-replicate codec must source
/// σ from an integer hyper-decoder, where the integer σ is exact and this snap
/// is unnecessary.
pub fn quantize_scale(sigma: f32, frac: u32) -> i64 {
    let s = (1u64 << frac) as f64;
    ((sigma.max(0.0) as f64) * s).floor() as i64
}

/// Integer `build_index`: bucket a fixed-point `σ` to a scale-table index, with
/// the same monotone semantics as the float `hyperprior::build_index` (index
/// rises with σ, clamped to `[0, len-1]`), but pure integer — **identical on
/// every machine** given the same `sigma_q`. `edges_q` comes from
/// [`fixed_scale_edges`]; `sigma_q` from an integer hyper-decoder (exact) or
/// [`quantize_scale`] (float source).
pub fn det_build_index(sigma_q: i64, edges_q: &[i64]) -> usize {
    let n = edges_q.len();
    debug_assert!(n >= 1, "scale table must be non-empty");
    let sigma_q = sigma_q.max(edges_q[0]); // never below the smallest tabulated scale
    let mut idx = n - 1;
    for &e in &edges_q[..n - 1] {
        if sigma_q <= e {
            idx -= 1;
        }
    }
    idx
}

/// Rebuild a latent element from its decoded symbol: `sym · 2^frac + base`, where
/// `base` is the element's integer mean (a conditional codec's μ) or its channel's
/// median (a factorized prior's), both at `2^frac`.
///
/// `None` when `frac > 62`, or when the exact value does not fit `i64`. The symbol
/// is read from the stream, and the bypass escape lets a stream carry any 32-bit
/// value, so an untrusted stream can ask for a value `i64` cannot hold; the caller
/// then refuses the stream. A plain `sym as i64 * one + base` would panic there in
/// an overflow-checked build and wrap without the checks, so the two builds would
/// decode the same bytes differently. Here both refuse. The arithmetic is `i128`,
/// where no step can overflow (`|sym · 2^62| < 2^94`); the Kani harness
/// `kani_dequantize_is_exact_or_refused` checks that for every input.
pub fn dequantize(sym: i32, frac: u32, base: i64) -> Option<i64> {
    if frac > 62 {
        return None;
    }
    i64::try_from(((sym as i128) << frac) + base as i128).ok()
}

/// The quantised CDF an encoder writes for a histogram, integer only: one table of
/// `hist.len() + 2` entries over the `hist.len()` symbols and the escape, the rANS
/// decoder's layout (`extensions/integer-payloads.md` §2; the escape is the last
/// symbol, `cdf_len − 2`). Every symbol, the escape included, gets frequency at least
/// 1, and the frequencies sum to exactly `2^precision`: the escape gets 1, each
/// symbol `1 + ⌊R·h_i / H⌋` with `R = 2^precision − hist.len() − 1` and `H` the
/// histogram's total, and the `R − Σ⌊·⌋` left over go one each to the symbols with
/// the largest remainders `R·h_i mod H`, ties to the lower index. An all-zero
/// histogram counts every symbol once. So the table is strictly increasing, from 0 to
/// `2^precision`, and the same histogram gives the same table on every machine.
///
/// `None` when the histogram is empty, when `precision` is outside `1..=16` (the
/// decoder's precision is 16), or when the symbols and the escape cannot each have
/// frequency 1 (`hist.len() + 1 > 2^precision`).
pub fn hist_to_cdf(hist: &[u64], precision: u32) -> Option<Vec<i32>> {
    if hist.is_empty() || !(1..=16).contains(&precision) {
        return None;
    }
    let total: u128 = 1u128 << precision;
    let n = hist.len() as u128;
    if n + 1 > total {
        return None;
    }
    let r = total - n - 1;
    let all_zero = hist.iter().all(|&h| h == 0);
    let weight = |h: u64| -> u128 { if all_zero { 1 } else { h as u128 } };
    let h_total: u128 = hist.iter().map(|&h| weight(h)).sum();
    // Room for the escape too, so pushing it below never regrows the vector: the fit
    // holds one frequency per symbol and the escape, never twice that.
    let mut freq: Vec<u128> = Vec::with_capacity(hist.len() + 1);
    freq.extend(hist.iter().map(|&h| 1 + r * weight(h) / h_total));
    let given: u128 = freq.iter().map(|f| f - 1).sum();
    let mut order: Vec<usize> = (0..hist.len()).collect();
    // Largest remainder first; equal remainders to the lower index (a stable sort).
    order.sort_by_key(|&i| std::cmp::Reverse(r * weight(hist[i]) % h_total));
    for &i in order.iter().take((r - given) as usize) {
        freq[i] += 1;
    }
    freq.push(1); // the escape
    let mut cdf = Vec::with_capacity(freq.len() + 1);
    let mut acc: u128 = 0;
    cdf.push(0i32);
    for f in freq {
        acc += f;
        cdf.push(acc as i32);
    }
    debug_assert_eq!(acc, total);
    Some(cdf)
}

// ─────────────────────────────────────────────────────────────────────
// Kani proof harnesses — the soundness hinge, checked rather than asserted
// ─────────────────────────────────────────────────────────────────────
//
// SPEC §7 makes this function load-bearing: a bucket that flips selects a
// different CDF and desyncs the remainder of the stream. Two things therefore
// have to hold, and both are proved here rather than argued in prose.
//
//   1. The index addresses a real scale-table entry. An out-of-range bucket is
//      either an out-of-bounds panic downstream or, worse, a silently wrong CDF.
//   2. The bucketing is *monotone* in σ. The doc comment claims "index rises
//      with σ"; a non-monotone bucketing would mis-order the table without ever
//      panicking, which no amount of round-trip testing on well-formed vectors
//      would reliably surface.
//
// Note the second is proved for an ARBITRARY edge table, not only an ascending
// one — the property follows from counting, not from the table being sorted, so
// the proof does not depend on a precondition a caller might not maintain.

/// The bucket always addresses a real entry of the scale table.
#[cfg(kani)]
#[kani::proof]
#[kani::unwind(6)]
fn kani_det_build_index_in_range() {
    let edges: [i64; 4] = kani::any();
    let sigma: i64 = kani::any();
    let idx = det_build_index(sigma, &edges);
    kani::assert(idx < edges.len(), "bucket index addresses a real scale-table entry");
    kani::cover!(true, "the bucketing is reachable");
}

/// **Monotonicity**: a larger σ never selects a smaller bucket.
///
/// This is the semantic property the integer bucketing exists to guarantee. It
/// holds by counting — a larger σ satisfies `σ <= e` for no more edges than a
/// smaller one does — so it is proved without assuming the table is ascending.
#[cfg(kani)]
#[kani::proof]
#[kani::unwind(6)]
fn kani_det_build_index_is_monotone_in_sigma() {
    let edges: [i64; 4] = kani::any();
    let a: i64 = kani::any();
    let b: i64 = kani::any();
    kani::assume(a <= b);
    kani::assert(
        det_build_index(a, &edges) <= det_build_index(b, &edges),
        "a larger sigma never selects a smaller bucket",
    );
    kani::cover!(true, "both bucketings are reachable");
}

/// The single-entry table degenerates correctly: one bucket, always index 0.
#[cfg(kani)]
#[kani::proof]
#[kani::unwind(4)]
fn kani_det_build_index_singleton_table() {
    let edges: [i64; 1] = kani::any();
    let sigma: i64 = kani::any();
    kani::assert(det_build_index(sigma, &edges) == 0, "a one-entry table has one bucket");
}

/// `dequantize` over every symbol, scale and base: it never panics, and it returns
/// the exact `sym · 2^frac + base` precisely when that value fits `i64` (and
/// `frac ≤ 62`), `None` otherwise. The expected value is formed by multiplication
/// by a power of two, not by the shift the function uses, so the harness states the
/// specification rather than restating the code.
#[cfg(kani)]
#[kani::proof]
fn kani_dequantize_is_exact_or_refused() {
    let sym: i32 = kani::any();
    let frac: u32 = kani::any();
    let base: i64 = kani::any();
    let got = dequantize(sym, frac, base);
    if frac <= 62 {
        let exact = sym as i128 * (1i128 << frac) + base as i128;
        let fits = exact >= i64::MIN as i128 && exact <= i64::MAX as i128;
        kani::assert(got.is_some() == fits, "refused exactly when the value leaves i64");
        if let Some(v) = got {
            kani::assert(v as i128 == exact, "the value returned is the exact one");
        }
    } else {
        kani::assert(got.is_none(), "a scale above 2^62 is refused");
    }
    kani::cover!(got.is_some(), "an in-range value is returned");
    kani::cover!(got.is_none() && frac <= 62, "an out-of-range value is refused");
}

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

    /// The float `build_index` the conditional family uses today — copied here
    /// (it is `pub(crate)` in the ORT-gated `neural` module) to demonstrate the
    /// hazard this module removes. Byte-for-byte the same logic as
    /// `neural::hyperprior::build_index`.
    fn float_build_index(sigma: f32, scale_table: &[f32]) -> usize {
        let n = scale_table.len();
        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
    }

    /// CompressAI's default 64 log-spaced scale buckets, 0.11 … 256.
    fn scale_table() -> Vec<f32> {
        let (lo, hi, n) = (0.11f32, 256.0f32, 64usize);
        let (ll, lh) = (lo.ln(), hi.ln());
        (0..n).map(|i| (ll + (lh - ll) * i as f32 / (n - 1) as f32).exp()).collect()
    }

    const FRAC: u32 = 16; // Q16: grid step ~1.5e-5

    /// `hist_to_cdf`'s table: total exactly 2^16, every frequency at least 1,
    /// strictly increasing, the escape last with frequency 1, the same table for the
    /// same histogram, and frequencies ordered as the counts are. Includes histograms
    /// far larger than the total, all-zero and one-symbol histograms, and a ladder
    /// where the remainders tie.
    #[test]
    fn hist_to_cdf_properties() {
        let cases: Vec<Vec<u64>> = vec![
            vec![0],
            vec![5],
            vec![0, 0, 0],
            vec![1, 1, 1, 1, 1, 1, 1],
            vec![1_000_000_000_000, 1, 0, 3, u64::MAX / 4],
            (0..4095u64).map(|i| (i * 2654435761) % 1000).collect(),
            (0..300u64).map(|i| if i == 150 { 1 << 40 } else { 0 }).collect(),
        ];
        for hist in &cases {
            let cdf = hist_to_cdf(hist, 16).expect("a table");
            assert_eq!(cdf.len(), hist.len() + 2);
            assert_eq!(cdf[0], 0);
            assert_eq!(*cdf.last().unwrap(), 1 << 16, "the total is exactly 2^16");
            assert!(cdf.windows(2).all(|w| w[1] > w[0]), "every frequency is at least 1");
            assert_eq!(cdf[cdf.len() - 1] - cdf[cdf.len() - 2], 1, "the escape is the last bin, frequency 1");
            assert_eq!(hist_to_cdf(hist, 16).unwrap(), cdf, "deterministic");
            let f: Vec<i32> = cdf.windows(2).map(|w| w[1] - w[0]).collect();
            for i in 0..hist.len() {
                for j in 0..hist.len() {
                    if hist[i] > hist[j] {
                        assert!(f[i] >= f[j], "a larger count never gets a smaller frequency");
                    }
                }
            }
        }
        // Equal counts: the leftover goes to the lower indices.
        let cdf = hist_to_cdf(&[1, 1, 1], 16).unwrap();
        let f: Vec<i32> = cdf.windows(2).map(|w| w[1] - w[0]).collect();
        assert_eq!(f, vec![21845, 21845, 21845, 1]);
        assert_eq!(hist_to_cdf(&[1, 1, 1, 1], 16).unwrap()[1..].iter().zip(&cdf_ladder()).filter(|(a, b)| a != b).count(), 0);
        // Refusals: empty, a precision the decoder does not use, too many symbols.
        assert!(hist_to_cdf(&[], 16).is_none());
        assert!(hist_to_cdf(&[1], 0).is_none());
        assert!(hist_to_cdf(&[1], 17).is_none());
        assert!(hist_to_cdf(&vec![1; 65536], 16).is_none());
        assert!(hist_to_cdf(&vec![1; 65535], 16).is_some());
    }

    /// Four equal counts: R = 65531, 16382 each with 3 left over, to indices 0, 1, 2.
    fn cdf_ladder() -> Vec<i32> {
        vec![16384, 32768, 49152, 65535, 65536]
    }

    #[test]
    fn integer_bucketing_matches_float_in_bucket_interiors() {
        let st = scale_table();
        let edges = fixed_scale_edges(&st, FRAC);
        // Sample the geometric midpoint of every adjacent bucket pair — points
        // far (≫ one grid step) from any edge, where quantization is exact.
        for w in st.windows(2) {
            let mid = (w[0] * w[1]).sqrt();
            assert_eq!(
                det_build_index(quantize_scale(mid, FRAC), &edges),
                float_build_index(mid, &st),
                "interior σ={mid} must bucket identically to the float path",
            );
        }
    }

    #[test]
    fn float_path_desyncs_on_cross_platform_drift() {
        // THE hazard the contract names: two platforms' float hyper-decoders
        // produce σ for the same element differing by ~1e-5 (real ONNX drift).
        // At a bucket edge the float bucketing DISAGREES → the stream desyncs.
        let st = scale_table();
        let edge = st[40];
        let drift = edge * 1e-5;
        let sigma_a = edge - drift; // platform A
        let sigma_b = edge + drift; // platform B
        assert_ne!(
            float_build_index(sigma_a, &st),
            float_build_index(sigma_b, &st),
            "float build_index should flip the bucket across the edge (the desync)",
        );
    }

    #[test]
    fn integer_source_makes_the_bucket_platform_invariant() {
        // The fix: an INTEGER hyper-decoder emits the identical integer σ on
        // every platform (integer arithmetic has no drift), so there is no
        // per-platform σ divergence to flip the bucket. Model that single
        // deterministic σ_q and bucket it: the result is the same on A and B.
        let st = scale_table();
        let edges = fixed_scale_edges(&st, FRAC);
        let sigma_q = 37_000_i64; // whatever the integer hyper-decoder emitted
        let on_a = det_build_index(sigma_q, &edges);
        let on_b = det_build_index(sigma_q, &edges); // identical integer input on B
        assert_eq!(on_a, on_b, "integer bucketing is platform-invariant by construction");
        // And it is monotone: a larger integer σ never buckets lower.
        let mut prev = 0;
        for q in (edges[0]..*edges.last().unwrap()).step_by(997) {
            let idx = det_build_index(q, &edges);
            assert!(idx >= prev, "bucketing must be monotone in σ");
            prev = idx;
        }
    }

    #[test]
    fn clamps_below_the_smallest_scale() {
        let st = scale_table();
        let edges = fixed_scale_edges(&st, FRAC);
        assert_eq!(det_build_index(i64::MIN / 2, &edges), 0);
        assert_eq!(det_build_index(i64::MAX / 2, &edges), edges.len() - 1);
    }

    #[test]
    fn dequantize_is_the_fixed_point_latent_or_a_refusal() {
        // In range it is exactly the expression the conditional decoders use.
        for (sym, base) in [(0, 0), (5, -3000), (-7, 12_345), (i32::MAX, 0), (i32::MIN, 0)] {
            assert_eq!(dequantize(sym, 12, base), Some(sym as i64 * (1 << 12) + base));
        }
        // At the edge of i64: the largest value that fits, and one past it.
        assert_eq!(dequantize(1, 62, i64::MAX - (1 << 62)), Some(i64::MAX));
        assert_eq!(dequantize(1, 62, i64::MAX - (1 << 62) + 1), None);
        assert_eq!(dequantize(-2, 62, 0), Some(i64::MIN));
        assert_eq!(dequantize(-2, 62, -1), None);
        // A stream-supplied symbol that no latent at this scale can be.
        assert_eq!(dequantize(i32::MAX, 40, 0), None);
        // Scales beyond 2^62 are refused rather than shifted out of range.
        assert_eq!(dequantize(0, 63, 0), None);
        assert_eq!(dequantize(1, u32::MAX, 0), None);
    }
}