//! The predictors of a residual refinement (`WIR1`, `extensions/staged-delivery.md`
//! Appendix A.1): what the stage below predicts for each element of the next stage.
//!
//! - **identity**: the element of the stage below at the same position;
//! - **up2**: a picture at twice the resolution, by the half-pel arithmetic of
//!   `int_video` (`(a + b + 1) >> 1`, `(a + b + c + d + 2) >> 2`), odd sizes included;
//! - **rate**: PCM at another sample rate, by linear interpolation rounded half away
//!   from zero (`rdiv`), channels interleaved.
//!
//! Every value is an integer of at most 64 bits, and every index is in bounds for the
//! shapes the caller checked (the Kani harnesses below, run by `wai-verify`, check
//! `up2_at`'s bounds, `rdiv` against its definition and `rate_index`'s indices). This file is
//! self-contained (no `crate::` or `super::` path, `core` only) so the model checker
//! and the cross-architecture byte-equality crate include it as it is.
//!
//! A decode predicts through the integer operator set's E11 and E12
//! (`staged::wir1::ops::{up2_prediction, rate_prediction}`, `wai_ops::stage`), whose
//! harnesses are in `wai-ops`. The predictors here are the direct decode's, test-only
//! (`legacy_int`), the reference the differential holds those to; [`rate_len`], the
//! rate predictor's geometry, is the parser's and the sink's.

/// `n / d` rounded to the nearest integer, halves away from zero:
/// `sign(n)·⌊(2|n| + d) / 2d⌋`, for `d ≥ 1`, computed as the quotient `q` and remainder
/// `r` of `|n|` by `d`, plus one where `2r ≥ d`. In `u64`, where no step can overflow
/// (`r < d`, and `q < 2^62` whenever `d ≥ 2`), so it is total for every `n` and every
/// `d ≥ 1`; `d < 1` gives 0.
#[cfg(any(test, kani, feature = "legacy_int"))]
pub fn rdiv(n: i64, d: i64) -> i64 {
    if d < 1 {
        return 0;
    }
    let (m, d) = (n.unsigned_abs(), d as u64);
    let (q, r) = (m / d, m % d);
    let mag = q + (r >= d - r) as u64;
    // |result| ≤ 2^63, and 2^63 only for n = i64::MIN, d = 1, whose result is n.
    if n < 0 { (mag as i64).wrapping_neg() } else { mag as i64 }
}

/// The up2 prediction of element `(y, x, ch)` of an `out_h × out_w` picture from
/// `prev`, the `h2 × w2 × c` picture below it (rows, then columns, then channels),
/// where `h2 = ⌈out_h/2⌉` and `w2 = ⌈out_w/2⌉`. A neighbour past the last row or
/// column is that row or column (`min(Y, h2 − 1)`, `min(X, w2 − 1)`).
///
/// The caller holds `prev.len() == h2·w2·c`, `y < 2·h2`, `x < 2·w2` and `ch < c`;
/// anything else gives 0 rather than reading out of bounds.
#[cfg(any(test, kani, feature = "legacy_int"))]
pub fn up2_at(prev: &[u8], h2: usize, w2: usize, c: usize, y: usize, x: usize, ch: usize) -> i64 {
    if h2 == 0 || w2 == 0 || c == 0 || ch >= c || y >= 2 * h2 || x >= 2 * w2 {
        return 0;
    }
    let Some(n) = h2.checked_mul(w2).and_then(|v| v.checked_mul(c)) else { return 0 };
    if prev.len() != n {
        return 0;
    }
    let p = |yy: usize, xx: usize| -> i64 { prev[(yy.min(h2 - 1) * w2 + xx.min(w2 - 1)) * c + ch] as i64 };
    let (iy, fy, ix, fx) = (y >> 1, y & 1, x >> 1, x & 1);
    match (fy, fx) {
        (0, 0) => p(iy, ix),
        (0, _) => (p(iy, ix) + p(iy, ix + 1) + 1) >> 1,
        (_, 0) => (p(iy, ix) + p(iy + 1, ix) + 1) >> 1,
        _ => (p(iy, ix) + p(iy, ix + 1) + p(iy + 1, ix) + p(iy + 1, ix + 1) + 2) >> 2,
    }
}

/// The source indices the rate prediction of output sample `i` reads, and its
/// interpolation phase: `s = i·den`, `lo = s div num`, `ρ = s mod num`,
/// `hi = min(lo + 1, n − 1)`, for `n` source samples per channel. `None` when
/// `num`, `den` or `n` is 0, or when `lo ≥ n` (an `i` past `⌊n·num/den⌋ − 1`).
#[cfg(any(test, kani, feature = "legacy_int"))]
pub fn rate_index(i: u64, n: u64, num: u64, den: u64) -> Option<(u64, u64, u64)> {
    if num == 0 || den == 0 || n == 0 {
        return None;
    }
    // Below the length `rate_len` gives, i·den < n·num, which is under 2^58 for the
    // shapes a WIR1 payload allows (n ≤ 2^26, num < 2^32): checked, and None past u64.
    let s = i.checked_mul(den)?;
    let (lo, rho) = (s / num, s % num);
    if lo >= n {
        return None;
    }
    Some((lo, (lo + 1).min(n - 1), rho))
}

/// The rate prediction of channel `ch` of output sample `i` from `prev`, `n` samples of
/// `c` interleaved channels at the source rate `den`, for output at the rate `num`:
/// `rdiv(a·num + (b − a)·ρ, num)` with `a = prev[lo·c + ch]`, `b = prev[hi·c + ch]`
/// ([`rate_index`]). This is `a + (b − a)·ρ/num` rounded half away from zero, the
/// linear interpolation at source position `i·den/num`, in exact integer arithmetic.
/// 0 where the indices are out of range (the caller checked the shapes).
#[cfg(any(test, kani, feature = "legacy_int"))]
pub fn rate_at(prev: &[i16], c: usize, num: u32, den: u32, i: usize, ch: usize) -> i64 {
    if c == 0 || ch >= c || prev.len() % c != 0 {
        return 0;
    }
    let n = (prev.len() / c) as u64;
    let Some((lo, hi, rho)) = rate_index(i as u64, n, num as u64, den as u64) else { return 0 };
    let a = prev[lo as usize * c + ch] as i64;
    let b = prev[hi as usize * c + ch] as i64;
    rdiv(a * num as i64 + (b - a) * rho as i64, num as i64)
}

/// The number of output samples per channel the rate predictor gives from `n` source
/// samples: `⌊n·num/den⌋`, or `None` for `den` 0 or `n·num` past `u64` (a WIR1 unit
/// holds at most 2^26 samples, so `n·num` is under 2^58 there).
pub fn rate_len(n: u64, num: u32, den: u32) -> Option<u64> {
    if den == 0 {
        return None;
    }
    Some(n.checked_mul(num as u64)? / den as u64)
}

// ─────────────────────────────────────────────────────────────────────
// Kani proof harnesses (wai-verify includes this file)
// ─────────────────────────────────────────────────────────────────────

/// `up2_bounded`: on every picture of up to 2×2 pixels and at every output position
/// of the 2× picture (odd output sizes, which leave the last row or column unused,
/// included), the prediction lies between the least and the greatest of the
/// neighbours it reads, so the half-pel rounding never leaves their range.
#[cfg(kani)]
#[kani::proof]
#[kani::unwind(5)]
fn up2_bounded() {
    let h2: usize = kani::any();
    let w2: usize = kani::any();
    kani::assume((1..=2).contains(&h2) && (1..=2).contains(&w2));
    let prev: [u8; 4] = kani::any();
    let y: usize = kani::any();
    let x: usize = kani::any();
    kani::assume(y < 2 * h2 && x < 2 * w2);
    let v = up2_at(&prev[..h2 * w2], h2, w2, 1, y, x, 0);
    let p = |yy: usize, xx: usize| prev[yy.min(h2 - 1) * w2 + xx.min(w2 - 1)] as i64;
    let (iy, ix) = (y >> 1, x >> 1);
    let ns = [p(iy, ix), p(iy, ix + 1), p(iy + 1, ix), p(iy + 1, ix + 1)];
    let (mut lo, mut hi) = (ns[0], ns[0]);
    for n in ns {
        lo = lo.min(n);
        hi = hi.max(n);
    }
    kani::assert(lo <= v && v <= hi, "the prediction lies between its neighbours");
    kani::cover!(y & 1 == 1 && x & 1 == 1, "the four-neighbour case is reached");
}

/// `rdiv_matches_spec`: `rdiv(n, d)` is `|n|/d` rounded half away from zero with the
/// sign of `n`, stated without division: its magnitude `k` is the one integer with
/// `(2k − 1)·d ≤ 2|n| < (2k + 1)·d` (for `k = 0`, `2|n| < d`). For `|n| < 2^12` and `d`
/// in `1..=2^6`: the obligation is nonlinear, which the solver discharges only over a
/// window; the unit tests cover `i64`'s extremes.
#[cfg(kani)]
#[kani::proof]
fn rdiv_matches_spec() {
    let n: i64 = kani::any();
    let d: i64 = kani::any();
    kani::assume(n > -(1i64 << 12) && n < (1i64 << 12));
    kani::assume(d >= 1 && d <= (1i64 << 6));
    let v = rdiv(n, d);
    kani::assert(v == 0 || (v < 0) == (n < 0), "the sign is n's");
    let (k, m, d) = (v.unsigned_abs(), n.unsigned_abs(), d as u64);
    kani::assert(2 * m < (2 * k + 1) * d, "below the next half");
    kani::assert(k == 0 || (2 * k - 1) * d <= 2 * m, "at or above the half below");
    kani::cover!(k >= 1 && (2 * k - 1) * d == 2 * m, "a half is reached");
}

/// `rate_index_in_bounds`: for every output sample `i` the rate predictor forms
/// (`i < ⌊n·num/den⌋`), with `n` source samples of `c` channels, both indices it reads
/// are samples of the source (`lo < n`, `hi < n`), so `lo·c + ch < n·c` for every
/// channel `ch < c`; `hi` is `min(lo + 1, n − 1)` and the phase is below `num`. For
/// `n`, `num` and `den` up to 32 and up to 4 channels: the obligation
/// (`i < ⌊n·num/den⌋ ⇒ ⌊i·den/num⌋ < n`) is nonlinear, which the solver discharges
/// only over a narrow window; the unit tests and the corpus cover larger rates.
#[cfg(kani)]
#[kani::proof]
fn rate_index_in_bounds() {
    let n: u64 = kani::any();
    let num: u32 = kani::any();
    let den: u32 = kani::any();
    let i: u64 = kani::any();
    kani::assume((1..=32).contains(&n) && (1..=32).contains(&num) && (1..=32).contains(&den));
    let len = rate_len(n, num, den);
    kani::assert(len.is_some(), "the output length fits u64");
    kani::assume(i < len.unwrap_or(0));
    let ix = rate_index(i, n, num as u64, den as u64);
    kani::assert(ix.is_some(), "every output sample has its indices");
    if let Some((lo, hi, rho)) = ix {
        kani::assert(lo < n && hi < n, "both indices are source samples");
        kani::assert(hi == (lo + 1).min(n - 1), "hi is the next sample, held at the end");
        kani::assert(rho < num as u64, "the phase is below num");
        let c: u64 = kani::any();
        let ch: u64 = kani::any();
        kani::assume((1..=4).contains(&c) && ch < c);
        kani::assert(lo * c + ch < n * c && hi * c + ch < n * c, "every channel's index is in the buffer");
    }
    kani::cover!(num > den, "an upsampling is reached");
    kani::cover!(num < den, "a downsampling is reached");
}

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

    /// `rdiv` rounds halves away from zero, on both signs.
    #[test]
    fn rdiv_rounds_halves_away_from_zero() {
        for (n, d, want) in [(5, 2, 3), (-5, 2, -3), (4, 2, 2), (-4, 2, -2), (7, 3, 2), (-7, 3, -2), (8, 3, 3),
                             (-8, 3, -3), (1, 2, 1), (-1, 2, -1), (0, 5, 0), (i64::MAX, 1, i64::MAX), (i64::MIN + 1, 1, i64::MIN + 1)] {
            assert_eq!(rdiv(n, d), want, "rdiv({n}, {d})");
        }
        assert_eq!(rdiv(3, 0), 0);
    }

    /// up2 on a 2×2 picture to 3×3 (odd), at every position: the half-pel
    /// arithmetic, with the last row and column repeated.
    #[test]
    fn up2_matches_the_half_pel_rule_on_odd_sizes() {
        let prev = [10u8, 20, 30, 41];
        let want = [[10, 15, 20], [20, 25, 31], [30, 36, 41]];
        for (y, row) in want.iter().enumerate() {
            for (x, &v) in row.iter().enumerate() {
                assert_eq!(up2_at(&prev, 2, 2, 1, y, x, 0), v, "({y},{x})");
            }
        }
        // Past the picture, a wrong buffer length, or a channel out of range: 0, no panic.
        assert_eq!(up2_at(&prev, 2, 2, 1, 4, 0, 0), 0);
        assert_eq!(up2_at(&prev[..3], 2, 2, 1, 0, 0, 0), 0);
        assert_eq!(up2_at(&prev, 2, 2, 1, 0, 0, 1), 0);
    }
}
