//! Quality per prefix (`extensions/staged-measurement.md` §4), in integers.
//!
//! A record states quality as integers that every machine computes alike: sums of
//! squared differences, the reference's energy, a largest difference, and a structural
//! similarity in Q32 fixed point. Decibel figures are derived from those integers for
//! display only ([`psnr_db_display`], [`snr_db_display`], [`ssim_display`]), and the
//! only floating-point arithmetic in this module is in those three functions.
//!
//! **Structural similarity** (§4.2). Luma is `Y = (77·R + 150·G + 29·B + 128) >> 8`.
//! Windows are `ww = min(8, W)` by `wh = min(8, H)` pixels. Their left edges are at
//! every multiple of 4 up to `W − ww`, and at `W − ww` itself when it is not one; their
//! top edges likewise up to `H − wh`. A window is at every pair of the two, so every
//! pixel is in one. For a window of `n = ww·wh` pixels, with `Sx = Σx`, `Sy = Σy`,
//! `Sxx = Σx²`, `Syy = Σy²` and `Sxy = Σxy` over the reference's luma `x` and the
//! test's luma `y`:
//!
//! ```text
//! A1 = 20000·Sx·Sy + 65025·n²
//! A2 = 20000·(n·Sxy − Sx·Sy) + 585225·n²
//! B1 = 10000·(Sx² + Sy²) + 65025·n²
//! B2 = 10000·((n·Sxx − Sx²) + (n·Syy − Sy²)) + 585225·n²
//! s_w = ⌊A1·A2·2³² / (B1·B2)⌋
//! ```
//!
//! and the figure is `⌊Σ s_w / windows⌋`. Both divisions floor toward −∞. This is the
//! structural similarity of Wang et al. (2004) with `K1 = 0.01`, `K2 = 0.03`, a peak of
//! 255, uniform windows and population statistics, with every factor multiplied out.
//!
//! **Why it fits `i128`.** A window holds `n ≤ 64` pixels of values `0..=255`, so
//! `0 ≤ Sx, Sy ≤ 16320` and `A1 ≤ 20000·16320² + 65025·64² < 5.33·10¹²`.
//! `n·Sxy − Sx·Sy` is `n²` times the window's covariance, whose magnitude is at most
//! the product of the two standard deviations, each at most 127.5, so
//! `|n·Sxy − Sx·Sy| ≤ 4096·16256.25 < 6.66·10⁷` and `|A2| < 1.34·10¹²`. Then
//! `|A1·A2·2³²| < 7.2·10²⁴·2³² < 2¹¹⁵`, well inside `i128`. `B1` and `B2` are each at
//! least their `n²` term, so `B1·B2 > 0`. `B1 ≤ 10000·2·16320² + 65025·64² < 5.33·10¹²`,
//! and `n·Sxx − Sx²` is `n²` times a variance of at most `127.5²`, so
//! `B2 ≤ 10000·2·4096·16256.25 + 585225·64² < 1.34·10¹²`; then
//! `B1·B2 < 7.2·10²⁴ < 2⁸³`. A similarity lies in `[−1, 1]`, so `|s_w| ≤ 2³²`, and a
//! sum over fewer than 2⁹⁰ windows fits as well: a picture has at most
//! `(⌊(W − 1)/4⌋ + 1)·(⌊(H − 1)/4⌋ + 1)` windows, never more than it has pixels.
//!
//! Several units are pooled (a clip, or every unit of a prefix) by summing each sum and
//! each window's `s_w` over the units and dividing once ([`image_pooled`]).
#![deny(clippy::float_arithmetic)]

/// An RGB8 picture: `h` rows of `w` pixels, `R, G, B` bytes, rows top to bottom.
#[derive(Clone, Copy, Debug)]
pub struct Rgb8<'a> {
    pub w: usize,
    pub h: usize,
    pub px: &'a [u8],
}

/// §4.2's integers for one picture, or for several pooled.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ImageQuality {
    /// Bytes compared: `3·W·H` per picture.
    pub n_samples: u64,
    /// `Σ (aᵢ − bᵢ)²` over the RGB bytes.
    pub sse_rgb: u128,
    /// The same sum over the luma planes.
    pub sse_y: u128,
    /// The mean structural similarity of the luma planes, in Q32 (`2³²` is 1).
    pub ssim_y_q32: i64,
    /// The windows the mean is over.
    pub windows: u64,
}

/// §4.3's integers for two sequences of `i16` samples.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct PcmQuality {
    pub n_samples: u64,
    /// `Σ (aᵢ − bᵢ)²`.
    pub sse: u128,
    /// `Σ aᵢ²` over the reference.
    pub ref_energy: u128,
}

/// §4.4's integers for two integer tensors.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct TensorQuality {
    pub n: u64,
    /// `Σ (aᵢ − bᵢ)²`.
    pub sse: u128,
    /// `max |aᵢ − bᵢ|`, 0 when equal.
    pub max_abs: u64,
}

/// Why two outputs have no quality figure.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum QualityError {
    /// Their sizes differ, or a picture's buffer is not `3·W·H` bytes.
    SizeMismatch,
    /// There is nothing to compare: no sample, or a picture with no pixel.
    Empty,
    /// A sum passed `u128` (only tensors of 64-bit values far apart reach it).
    Overflow,
}

impl std::fmt::Display for QualityError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            Self::SizeMismatch => "the two outputs differ in size",
            Self::Empty => "nothing to compare",
            Self::Overflow => "a sum passes 128 bits",
        })
    }
}

impl std::error::Error for QualityError {}

/// The integer BT.601 luma of RGB8 bytes: `(77·R + 150·G + 29·B + 128) >> 8` per
/// pixel. A trailing partial pixel is ignored.
pub fn luma_bt601_int(px: &[u8]) -> Vec<u8> {
    px.chunks_exact(3)
        .map(|p| ((77 * p[0] as u32 + 150 * p[1] as u32 + 29 * p[2] as u32 + 128) >> 8) as u8)
        .collect()
}

/// One picture's sums, with `Σ s_w` and its window count, before the mean is taken.
struct ImageSums {
    n: u64,
    sse_rgb: u128,
    sse_y: u128,
    ssim_sum: i128,
    windows: u64,
}

fn sq(a: u8, b: u8) -> u128 {
    let d = a.abs_diff(b) as u128;
    d * d
}

/// A window's offsets along an axis of `len` pixels: every multiple of 4 up to
/// `len − win`, then `len − win` itself when it is not one, so the last window ends at
/// the edge and every pixel is in a window.
fn starts(len: usize, win: usize) -> impl Iterator<Item = usize> {
    let last = len - win;
    (0..=last).step_by(4).chain((last % 4 != 0).then_some(last))
}

fn image_sums(r: Rgb8<'_>, t: Rgb8<'_>) -> Result<ImageSums, QualityError> {
    if r.w != t.w || r.h != t.h {
        return Err(QualityError::SizeMismatch);
    }
    if r.w == 0 || r.h == 0 {
        return Err(QualityError::Empty);
    }
    let len = r.w.checked_mul(r.h).and_then(|n| n.checked_mul(3)).ok_or(QualityError::SizeMismatch)?;
    if r.px.len() != len || t.px.len() != len {
        return Err(QualityError::SizeMismatch);
    }
    let sse_rgb = r.px.iter().zip(t.px).map(|(&a, &b)| sq(a, b)).sum();
    let (ya, yb) = (luma_bt601_int(r.px), luma_bt601_int(t.px));
    let sse_y = ya.iter().zip(&yb).map(|(&a, &b)| sq(a, b)).sum();
    let (w, h) = (r.w, r.h);
    let (ww, wh) = (w.min(8), h.min(8));
    let n = (ww * wh) as i128;
    let (c1, c2) = (65025 * n * n, 585225 * n * n);
    let mut ssim_sum: i128 = 0;
    let mut windows = 0u64;
    for y0 in starts(h, wh) {
        for x0 in starts(w, ww) {
            let (mut sx, mut sy, mut sxx, mut syy, mut sxy) = (0i128, 0i128, 0i128, 0i128, 0i128);
            for y in y0..y0 + wh {
                for x in x0..x0 + ww {
                    let (p, q) = (ya[y * w + x] as i128, yb[y * w + x] as i128);
                    sx += p;
                    sy += q;
                    sxx += p * p;
                    syy += q * q;
                    sxy += p * q;
                }
            }
            let a1 = 20000 * sx * sy + c1;
            let a2 = 20000 * (n * sxy - sx * sy) + c2;
            let b1 = 10000 * (sx * sx + sy * sy) + c1;
            let b2 = 10000 * ((n * sxx - sx * sx) + (n * syy - sy * sy)) + c2;
            ssim_sum += (a1 * a2 * (1i128 << 32)).div_euclid(b1 * b2);
            windows += 1;
        }
    }
    Ok(ImageSums { n: len as u64, sse_rgb, sse_y, ssim_sum, windows })
}

/// §4.2 for one picture against its reference.
pub fn image(reference: Rgb8<'_>, test: Rgb8<'_>) -> Result<ImageQuality, QualityError> {
    image_pooled(&[(reference, test)])
}

/// §4.2's sums pooled over pictures added one at a time (§4.5): every sum is summed
/// over the pictures, and the similarity is the mean of every window of every
/// picture, one division at the end. A pool holds no picture, so pooling a prefix's
/// units needs only one unit's output at a time.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct ImagePool {
    n: u64,
    sse_rgb: u128,
    sse_y: u128,
    ssim_sum: i128,
    windows: u64,
}

impl ImagePool {
    /// Add one picture against its reference; the pair must match in size.
    pub fn add(&mut self, reference: Rgb8<'_>, test: Rgb8<'_>) -> Result<(), QualityError> {
        let s = image_sums(reference, test)?;
        self.n += s.n;
        self.sse_rgb += s.sse_rgb;
        self.sse_y += s.sse_y;
        self.ssim_sum += s.ssim_sum;
        self.windows += s.windows;
        Ok(())
    }

    /// The pooled figure, or `None` when no picture was added.
    pub fn finish(&self) -> Option<ImageQuality> {
        // Every picture has at least one window (`ww ≤ W`, `wh ≤ H`), so a pool with a
        // picture has `windows ≥ 1`.
        (self.windows > 0).then(|| ImageQuality {
            n_samples: self.n,
            sse_rgb: self.sse_rgb,
            sse_y: self.sse_y,
            ssim_y_q32: self.ssim_sum.div_euclid(self.windows as i128) as i64,
            windows: self.windows,
        })
    }
}

/// §4.2 pooled over several pictures (a clip's frames, or a prefix's units): every
/// sum is summed over the pictures, and the similarity is the mean of every window
/// of every picture. Each pair must match in size; the pictures may differ.
pub fn image_pooled(pairs: &[(Rgb8<'_>, Rgb8<'_>)]) -> Result<ImageQuality, QualityError> {
    let mut pool = ImagePool::default();
    for &(r, t) in pairs {
        pool.add(r, t)?;
    }
    pool.finish().ok_or(QualityError::Empty)
}

/// §4.3 for two sequences of `i16` samples, channel-interleaved in output order.
/// Sequences of different lengths are refused.
pub fn pcm16(reference: &[i16], test: &[i16]) -> Result<PcmQuality, QualityError> {
    if reference.len() != test.len() {
        return Err(QualityError::SizeMismatch);
    }
    if reference.is_empty() {
        return Err(QualityError::Empty);
    }
    let mut q = PcmQuality { n_samples: reference.len() as u64, sse: 0, ref_energy: 0 };
    for (&a, &b) in reference.iter().zip(test) {
        let d = (a as i32 - b as i32).unsigned_abs() as u128;
        q.sse += d * d;
        let e = (a as i32).unsigned_abs() as u128;
        q.ref_energy += e * e;
    }
    Ok(q)
}

/// [`pcm16`] pooled over several pairs.
pub fn pcm16_pooled(pairs: &[(&[i16], &[i16])]) -> Result<PcmQuality, QualityError> {
    if pairs.is_empty() {
        return Err(QualityError::Empty);
    }
    let mut out = PcmQuality { n_samples: 0, sse: 0, ref_energy: 0 };
    for &(r, t) in pairs {
        let q = pcm16(r, t)?;
        out.n_samples += q.n_samples;
        out.sse += q.sse;
        out.ref_energy += q.ref_energy;
    }
    Ok(out)
}

/// §4.4 for two integer tensors in the capability's canonical order.
pub fn tensor_i64(reference: &[i64], test: &[i64]) -> Result<TensorQuality, QualityError> {
    tensor_pooled(&[(reference, test)])
}

/// [`tensor_i64`] pooled over several pairs: the sums are summed and `max_abs` is the
/// largest.
pub fn tensor_pooled(pairs: &[(&[i64], &[i64])]) -> Result<TensorQuality, QualityError> {
    if pairs.is_empty() {
        return Err(QualityError::Empty);
    }
    let mut out = TensorQuality { n: 0, sse: 0, max_abs: 0 };
    for &(r, t) in pairs {
        if r.len() != t.len() {
            return Err(QualityError::SizeMismatch);
        }
        if r.is_empty() {
            return Err(QualityError::Empty);
        }
        for (&a, &b) in r.iter().zip(t) {
            let d = (a as i128 - b as i128).unsigned_abs();
            // |a − b| < 2⁶⁴ for two i64, so it fits u64; its square fits u128.
            out.max_abs = out.max_abs.max(d as u64);
            out.sse = out.sse.checked_add(d * d).ok_or(QualityError::Overflow)?;
        }
        out.n += r.len() as u64;
    }
    Ok(out)
}

/// PSNR in dB for `n` 8-bit samples with squared error `sse`:
/// `10·log10(255²·n / sse)`. For display only (§4.6); `None` when `sse` is 0 (no
/// finite figure) or `n` is 0.
#[allow(clippy::float_arithmetic)]
pub fn psnr_db_display(sse: u128, n: u64) -> Option<f64> {
    if sse == 0 || n == 0 {
        return None;
    }
    Some(10.0 * (65025.0 * n as f64 / sse as f64).log10())
}

/// SNR in dB: `10·log10(ref_energy / sse)`. For display only (§4.6); `None` when
/// `sse` or `ref_energy` is 0.
#[allow(clippy::float_arithmetic)]
pub fn snr_db_display(ref_energy: u128, sse: u128) -> Option<f64> {
    if sse == 0 || ref_energy == 0 {
        return None;
    }
    Some(10.0 * (ref_energy as f64 / sse as f64).log10())
}

/// A Q32 similarity as a number in `[−1, 1]`, for display only.
#[allow(clippy::float_arithmetic)]
pub fn ssim_display(q32: i64) -> f64 {
    q32 as f64 / 4_294_967_296.0
}

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

    /// A deterministic test picture. The literals below were computed from it in Python
    /// twice, by `staged-measure-conformance/measure_verify.py`'s `image_parts` and by a
    /// rational-arithmetic structural similarity with its own window placement.
    fn picture(w: usize, h: usize, seed: usize, mul: usize) -> Vec<u8> {
        let mut out = Vec::with_capacity(w * h * 3);
        for y in 0..h {
            for x in 0..w {
                for c in 0..3 {
                    out.push(((x * mul + y * (mul + 3) + c * 29 + seed * 17 + ((x * y + seed) % 13) * 5) % 256) as u8);
                }
            }
        }
        out
    }

    fn q(w: usize, h: usize, a: &[u8], b: &[u8]) -> ImageQuality {
        image(Rgb8 { w, h, px: a }, Rgb8 { w, h, px: b }).unwrap()
    }

    fn iq(n_samples: u64, sse_rgb: u128, sse_y: u128, ssim_y_q32: i64, windows: u64) -> ImageQuality {
        ImageQuality { n_samples, sse_rgb, sse_y, ssim_y_q32, windows }
    }

    #[test]
    fn sse_literals() {
        // Hand values: (255² · 3) + (3² + 4²) + (1 + 1 + 9) = 195075 + 25 + 11.
        let a = [0, 0, 0, 255, 255, 255, 10, 20, 30, 200, 100, 50];
        let b = [0, 0, 0, 0, 0, 0, 13, 16, 30, 201, 99, 53];
        assert_eq!(q(2, 2, &a, &b), iq(12, 195111, 65026, 187211645, 1));
        assert_eq!(q(2, 2, &[7; 12], &[7; 12]), iq(12, 0, 0, 1 << 32, 1));
        assert_eq!(q(2, 2, &[255; 12], &[0; 12]), iq(12, 780300, 260100, 429453, 1));
    }

    #[test]
    fn ssim_identical_is_one() {
        for (w, h) in [(1, 1), (8, 8), (13, 5), (40, 24)] {
            let p = picture(w, h, 2, 9);
            assert_eq!(q(w, h, &p, &p).ssim_y_q32, 1 << 32, "{w}x{h}");
        }
    }

    #[test]
    fn ssim_extreme_windows() {
        // A full 8×8 window at the extremes: all white against all black, and a
        // checkerboard against its inverse (the most negative covariance a window has).
        assert_eq!(q(8, 8, &[255; 192], &[0; 192]), iq(192, 12484800, 4161600, 429453, 1));
        let cb: Vec<u8> = (0..64).flat_map(|i| [if (i % 8 + i / 8) % 2 == 1 { 255u8 } else { 0 }; 3]).collect();
        let inv: Vec<u8> = cb.iter().map(|v| 255 - v).collect();
        assert_eq!(q(8, 8, &cb, &inv), iq(192, 12484800, 4161600, -4279533196, 1));
    }

    #[test]
    fn ssim_small_planes() {
        // Planes narrower or shorter than a window use clipped windows: 3 × 8 down a
        // 3 × 100 plane at stride 4 (24 windows), and one 5 × 5 window.
        assert_eq!(q(3, 100, &picture(3, 100, 1, 7), &picture(3, 100, 2, 7)), iq(900, 4705154, 532902, 2841543333, 24));
        assert_eq!(q(5, 5, &picture(5, 5, 3, 11), &picture(5, 5, 3, 13)), iq(75, 65392, 2145, 4227018163, 1));
    }

    #[test]
    fn ssim_literals() {
        // (w, h, seed a, mul a, seed b, mul b) → n, sse_rgb, sse_y, ssim_y_q32, windows,
        // as the independent reference computes them.
        // 17 × 9, 9 × 33 and 11 × 11 have an edge-aligned last window on an axis.
        let cases: [((usize, usize, usize, usize, usize, usize), ImageQuality); 7] = [
            ((16, 16, 1, 5, 1, 6), iq(768, 3329152, 245154, 3542145100, 9)),
            ((17, 9, 2, 3, 5, 3), iq(459, 2191911, 500794, 1115121540, 8)),
            ((12, 20, 4, 9, 5, 9), iq(720, 4072752, 548256, 2404129353, 8)),
            ((32, 8, 6, 1, 7, 2), iq(768, 8646871, 893709, -424684373, 7)),
            ((9, 33, 8, 13, 9, 14), iq(891, 8244787, 1469842, 1291000068, 16)),
            ((40, 24, 3, 7, 3, 8), iq(2880, 19201376, 3078475, 2198626257, 45)),
            ((11, 11, 1, 5, 2, 7), iq(363, 1366484, 201849, 3172402069, 4)),
        ];
        for ((w, h, s1, m1, s2, m2), want) in cases {
            assert_eq!(q(w, h, &picture(w, h, s1, m1), &picture(w, h, s2, m2)), want, "{w}x{h}");
        }
    }

    #[test]
    fn pooled_is_one_mean_over_every_window() {
        let (a0, b0) = (picture(16, 16, 1, 5), picture(16, 16, 1, 6));
        let (a1, b1) = (picture(17, 9, 2, 3), picture(17, 9, 5, 3));
        let p = image_pooled(&[
            (Rgb8 { w: 16, h: 16, px: &a0 }, Rgb8 { w: 16, h: 16, px: &b0 }),
            (Rgb8 { w: 17, h: 9, px: &a1 }, Rgb8 { w: 17, h: 9, px: &b1 }),
        ])
        .unwrap();
        // ⌊(31879305901 + 8920972324) / 17⌋, not the mean of the two pictures' means.
        assert_eq!(p, iq(1227, 5521063, 745948, 2400016366, 17));
    }

    #[test]
    fn every_pixel_is_in_a_window() {
        // Offsets per axis: multiples of 4 up to `len − win`, then the edge-aligned one.
        assert_eq!(starts(11, 8).collect::<Vec<_>>(), [0, 3]);
        assert_eq!(starts(16, 8).collect::<Vec<_>>(), [0, 4, 8]);
        assert_eq!(starts(17, 8).collect::<Vec<_>>(), [0, 4, 8, 9]);
        assert_eq!(starts(5, 5).collect::<Vec<_>>(), [0]);
        for len in 1..=70 {
            let win = len.min(8);
            let mut seen = vec![false; len];
            let v: Vec<usize> = starts(len, win).collect();
            for &x in &v {
                seen[x..x + win].iter_mut().for_each(|s| *s = true);
            }
            assert!(seen.iter().all(|&s| s), "len {len}");
            // Never more offsets than the stated bound, and no offset twice.
            assert!(v.len() <= (len - 1) / 4 + 1 && v.windows(2).all(|w| w[0] < w[1]), "len {len}");
        }
        // An 11 × 11 picture whose only difference is its last row and column: with
        // windows at 0 alone it would score 1; the edge-aligned windows see it.
        let a = picture(11, 11, 4, 3);
        let mut b = a.clone();
        for i in 0..11 {
            for c in 0..3 {
                b[(10 * 11 + i) * 3 + c] ^= 0x40;
                b[(i * 11 + 10) * 3 + c] ^= 0x40;
            }
        }
        let got = q(11, 11, &a, &b);
        assert_eq!(got.windows, 4);
        assert!(got.ssim_y_q32 < 1 << 32);
    }

    #[test]
    fn sizes_are_checked() {
        let p = [0u8; 12];
        assert_eq!(image(Rgb8 { w: 2, h: 2, px: &p }, Rgb8 { w: 1, h: 4, px: &p }), Err(QualityError::SizeMismatch));
        assert_eq!(image(Rgb8 { w: 2, h: 2, px: &p }, Rgb8 { w: 2, h: 2, px: &p[..9] }), Err(QualityError::SizeMismatch));
        assert_eq!(image(Rgb8 { w: 0, h: 4, px: &[] }, Rgb8 { w: 0, h: 4, px: &[] }), Err(QualityError::Empty));
        assert_eq!(pcm16(&[1, 2], &[1]), Err(QualityError::SizeMismatch));
        assert_eq!(pcm16(&[], &[]), Err(QualityError::Empty));
        assert_eq!(tensor_i64(&[1], &[1, 2]), Err(QualityError::SizeMismatch));
    }

    #[test]
    fn pcm_and_tensor_literals() {
        // sse = 2² + 32767² + 65535²; ref_energy = 10² + 32767² + 32768².
        let q = pcm16(&[10, -32767, -32768], &[12, 0, 32767]).unwrap();
        assert_eq!(q, PcmQuality { n_samples: 3, sse: 4 + 1073676289 + 4294836225, ref_energy: 100 + 1073676289 + 1073741824 });
        // The widest i64 gap: |i64::MIN − i64::MAX| = 2⁶⁴ − 1, squared within u128.
        let t = tensor_i64(&[i64::MIN, 5, 7], &[i64::MAX, 5, 4]).unwrap();
        assert_eq!(t, TensorQuality { n: 3, sse: (u64::MAX as u128) * (u64::MAX as u128) + 9, max_abs: u64::MAX });
        // Two such gaps pass u128.
        assert_eq!(tensor_i64(&[i64::MIN, i64::MIN], &[i64::MAX, i64::MAX]), Err(QualityError::Overflow));
    }

    #[test]
    fn display_figures() {
        assert_eq!(psnr_db_display(0, 10), None);
        // 255²·1 / 65025 = 1 → 0 dB.
        assert_eq!(psnr_db_display(65025, 1), Some(0.0));
        assert_eq!(snr_db_display(7, 7), Some(0.0));
        assert_eq!(ssim_display(1 << 31), 0.5);
    }
}
