maroontree 0.1.8

/*
 * Copyright (c) Radzivon Bartoshyk 6/2026. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without modification,
 * are permitted provided that the following conditions are met:
 *
 * 1.  Redistributions of source code must retain the above copyright notice, this
 * list of conditions and the following disclaimer.
 *
 * 2.  Redistributions in binary form must reproduce the above copyright notice,
 * this list of conditions and the following disclaimer in the documentation
 * and/or other materials provided with the distribution.
 *
 * 3.  Neither the name of the copyright holder nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

use crate::av2::entropy::RangeEncoder;
use crate::av2::helpers::sum_sumsq_f32;
use crate::av2::quant::qstep;
use crate::util::FastRound;

pub(crate) const AQ_RES_LOG2: u8 = 2;

const AQ_MAX_SIGNALED: i32 = 6;

pub(crate) fn sb_activity(
    yp: &[f32],
    pw: usize,
    sb_y: usize,
    sb_x: usize,
    width: usize,
    height: usize,
) -> f32 {
    let h = height.saturating_sub(sb_y).min(64);
    let w = width.saturating_sub(sb_x).min(64);
    if h == 0 || w == 0 {
        return 0.0;
    }
    let mut sum = 0f32;
    let mut sum2 = 0f32;
    for r in 0..h {
        let base = (sb_y + r) * pw + sb_x;
        let (row_sum, row_sum2) = sum_sumsq_f32(&yp[base..base + w]);
        sum += row_sum;
        sum2 += row_sum2;
    }
    let n = (h * w) as f32;
    let mean = sum / n;
    let var = (sum2 / n - mean * mean).max(0.0);
    dirty_log1pf(var)
}

fn sb_subblock_variances(
    yp: &[f32],
    pw: usize,
    sb_y: usize,
    sb_x: usize,
    width: usize,
    height: usize,
    out: &mut [f32; 64],
) -> usize {
    let mut filled = 0usize;
    let mut acc = 0f32;
    for (by, row) in out.as_chunks_mut::<8>().0.iter_mut().take(8).enumerate() {
        for (bx, out) in row.iter_mut().enumerate() {
            let y0 = sb_y + by * 8;
            let x0 = sb_x + bx * 8;
            let h = height.saturating_sub(y0).min(8);
            let w = width.saturating_sub(x0).min(8);
            if h == 0 || w == 0 {
                *out = f32::NAN; // mark out-of-frame, patched below
                continue;
            }
            let mut sum = 0f32;
            let mut sum2 = 0f32;
            for r in 0..h {
                let base = (y0 + r) * pw + x0;
                let (row_sum, row_sum2) = sum_sumsq_f32(&yp[base..base + w]);
                sum += row_sum;
                sum2 += row_sum2;
            }
            let n = (h * w) as f32;
            let mean = sum / n;
            let var = (sum2 / n - mean * mean).max(0.0);
            *out = var;
            acc += var;
            filled += 1;
        }
    }
    if filled == 0 {
        out.iter_mut().for_each(|v| *v = 0.0);
        return 0;
    }
    // Patch out-of-frame subblocks with the in-frame mean (neutral for the octile).
    let mean = acc / filled as f32;
    for v in out.iter_mut() {
        if v.is_nan() {
            *v = mean;
        }
    }
    filled
}

/// Mean activity over all superblocks of a (padded) luma plane — the per-tile
/// reference used to center the AQ deltas so they are zero-mean.
pub(crate) fn tile_ref_activity(
    yp: &[f32],
    pw: usize,
    sb_rows: usize,
    sb_cols: usize,
    width: usize,
    height: usize,
) -> f32 {
    let mut sum = 0f32;
    let mut cnt = 0f32;
    for row in 0..sb_rows {
        for col in 0..sb_cols {
            sum += sb_activity(yp, pw, row * 64, col * 64, width, height);
            cnt += 1.0;
        }
    }
    if cnt > 0.0 { sum / cnt } else { 5.0 }
}

/// Scale a residual block by `s` (identity-fast for `s == 1.0`). Used to quantize
/// at a per-SB qstep via the linear projection: `levels = project(resid * (qstep_base
/// / qstep_sb))` equals quantizing at `qstep_sb`.
pub(crate) fn scale_resid(v: &[f32], s: f32) -> Vec<f32> {
    if s == 1.0 {
        v.to_vec()
    } else {
        v.iter().map(|&x| x * s).collect()
    }
}

/// Per-tile adaptive-quantization state. Construct once per `encode_*_core`
/// invocation (one tile). Inter-capable paths should use [`AqState::per_sb_probe`]
/// while choosing a mode and call [`AqState::per_sb`] only when the selected block
/// carries delta-Q syntax. A full-superblock `skip_txfm=1` block must not advance
/// `last_qidx`, because the decoder does not read delta-Q for that block.
/// One superblock's precomputed AQ decision. Wavefront blocker #1: the parallel
/// (diagonal-order) decide can't thread `per_sb`'s raster-serial `last_qidx`
/// accumulator, so a cheap serial pre-pass ([`AqState::precompute_grid`]) yields
/// this grid for the decide to read; `sig` feeds the serial emit's `delta_q`.
#[derive(Clone, Copy)]
#[allow(dead_code)]
pub(crate) struct AqCell {
    pub(crate) qs: i32,
    pub(crate) resid_scale: f32,
    pub(crate) qs_c: i32,
    pub(crate) resid_scale_c: f32,
    pub(crate) sig: i32,
    /// The committed qindex this SB codes at (`newq` for full-interior, the entering
    /// accumulator for edge/`!present`). Mirrors `current_qidx()` after this SB.
    pub(crate) qidx: i32,
    /// The accumulator qstep *entering* this SB, before it commits — i.e. what
    /// `aqs.current()` returns at the partition probe (the previous full-interior
    /// SB's `newq`). In a raster serial run this equals the prior cell's state; a
    /// wavefront decide reads it from the grid instead of a serial accumulator.
    pub(crate) qs_before: i32,
    pub(crate) resid_scale_before: f32,
}

// The dark-structured-detail protection config lives in the shared AQ module (used by
// both the AV1 and AV2 paths); re-exported here so `av2::DarkAq` (public API) resolves.
pub use crate::aq_common::DarkAq;
use crate::aq_common::dirty_log1pf;

pub(crate) struct AqState {
    present: bool,
    base_q: i32,
    qstep_base: i32,
    /// Frame-level chroma qindex delta (AVM base_uv_ac_delta_q); chroma qstep at an
    /// SB is `qstep((sb_qidx + uv_delta).clamp(1,255))`, matching decoder `get_q`.
    uv_delta: i32,
    ref_act: f32,
    /// Running qindex accumulator (decoder `ts.last_qidx`), reset to the frame
    /// base at each tile start.
    last_qidx: i32,
    /// Variance Boost selectivity octile (1..=8). Default 6 (SVT-AV1-PSY default).
    vb_octile: u8,
    /// Variance Boost strength multiplier (1.0 = nominal).
    vb_strength: f32,
    /// When true, only boost low-variance SBs (net-negative, spends bits). When false
    /// (default), also coarsen high-variance SBs to keep the rate roughly matched.
    vb_boost_only: bool,
    dark: DarkAq,
}

impl AqState {
    /// `present` = frame enables delta-Q; `qstep_base` = qstep at the frame base
    /// qindex; `ref_act` = tile mean activity (see [`tile_ref_activity`]).
    pub(crate) fn new(
        present: bool,
        base_q: i32,
        qstep_base: i32,
        ref_act: f32,
        uv_delta: i32,
    ) -> Self {
        AqState {
            present,
            base_q,
            qstep_base,
            uv_delta,
            ref_act,
            last_qidx: base_q,
            vb_octile: 6,
            vb_strength: 1.0,
            vb_boost_only: false,
            dark: DarkAq::default(),
        }
    }

    /// Enable/configure the dark-structured-detail protection term (see [`DarkAq`]).
    pub(crate) fn with_dark_aq(mut self, dark: DarkAq) -> Self {
        self.dark = dark;
        self
    }

    /// Chroma (qstep, resid_scale) at qindex `q`: decoder chroma dequant is
    /// `get_q(clamp(q + uv_delta, 1, 255))` (equal ac/dc), residual pre-scale keeps
    /// the shared basis calibrated at `qstep_base`.
    fn chroma_at(&self, q: i32) -> (i32, f32) {
        let qc = qstep((q + self.uv_delta).clamp(1, 255) as u32) as i32;
        (qc, self.qstep_base as f32 / qc as f32)
    }

    /// Override the Variance Boost knobs (octile selectivity, strength, boost-only).
    /// Returns `self` for chaining at the construction site.
    pub(crate) fn with_variance_boost(
        mut self,
        octile: u8,
        strength: f32,
        boost_only: bool,
    ) -> Self {
        self.vb_octile = octile.clamp(1, 8);
        // SS2-calibrated taper: boost pays at coarse q, is net-negative near-lossless.
        let taper = ((self.base_q as f32 - 30.0) / 40.0).clamp(0.0, 1.0);
        self.vb_strength = strength.max(0.0) * taper;
        self.vb_boost_only = boost_only;
        self
    }

    /// Extra qindex reduction (>= 0) for a dark, structured SB. 0 when disabled, out of
    /// the gated quality range, or the SB carries no cross-scale structure. Delegates to
    /// the shared implementation; the AV2 luma plane is native-depth `f32`, passed at
    /// `scale = 1.0` (no bit-depth normalization).
    fn dark_protection(
        &self,
        yp: &[f32],
        pw: usize,
        sb_y: usize,
        sb_x: usize,
        width: usize,
        height: usize,
    ) -> i32 {
        crate::aq_common::dark_protection(
            &self.dark,
            self.base_q,
            yp,
            pw,
            sb_y,
            sb_x,
            width,
            height,
            1.0,
        )
    }

    /// The SB's target qindex: the flat-region variance boost combined with the
    /// dark-structured-detail protection. `protection = max(flat_boost, dark_boost)`;
    /// the coarsen (positive) side of the variance boost is kept only when neither
    /// protection signal fires. Shared by `per_sb`, `per_sb_probe` and `precompute_grid`
    /// so they stay bit-exact.
    fn sb_target(
        &self,
        yp: &[f32],
        pw: usize,
        sb_y: usize,
        sb_x: usize,
        width: usize,
        height: usize,
    ) -> i32 {
        let mut subvars = [0f32; 64];
        let filled = sb_subblock_variances(yp, pw, sb_y, sb_x, width, height, &mut subvars);
        if filled == 0 {
            return self.base_q;
        }
        let picked = crate::aq_common::sb_octile_variance(&mut subvars, self.vb_octile);
        let vb_delta = crate::aq_common::variance_boost_delta(
            picked,
            self.ref_act,
            self.vb_strength,
            self.vb_boost_only,
        );
        let dark = self.dark_protection(yp, pw, sb_y, sb_x, width, height);
        let flat_boost = (-vb_delta).max(0);
        let protection = flat_boost.max(dark);
        let delta = if protection > 0 {
            -protection
        } else {
            vb_delta
        };
        (self.base_q + delta).clamp(1, 255)
    }

    pub(crate) fn current_qidx(&self) -> i32 {
        self.last_qidx
    }

    #[allow(clippy::too_many_arguments)]
    pub(crate) fn per_sb_probe(
        &self,
        yp: &[f32],
        pw: usize,
        sb_y: usize,
        sb_x: usize,
        width: usize,
        height: usize,
    ) -> (i32, f32, i32, f32) {
        if !self.present {
            let (qc, rc) = self.chroma_at(self.base_q);
            return (self.qstep_base, 1.0, qc, rc);
        }
        let target = self.sb_target(yp, pw, sb_y, sb_x, width, height);
        let step = 1i32 << AQ_RES_LOG2;
        let sig = (((target - self.last_qidx) as f32) / step as f32)
            .fast_round()
            .clamp(-(AQ_MAX_SIGNALED as f32), AQ_MAX_SIGNALED as f32) as i32;
        let newq = (self.last_qidx + sig * step).clamp(1, 255);
        let qs = qstep(newq as u32) as i32;
        let (qc, rc) = self.chroma_at(newq);
        (qs, self.qstep_base as f32 / qs as f32, qc, rc)
    }

    /// Current accumulated (qstep, resid_scale) without signaling — for SBs that
    /// emit delta 0 but must code at the decoder's accumulated qindex.
    pub(crate) fn current(&self) -> (i32, f32, i32, f32) {
        if !self.present {
            let (qc, rc) = self.chroma_at(self.base_q);
            return (self.qstep_base, 1.0, qc, rc);
        }
        let qs = qstep(self.last_qidx as u32) as i32;
        let (qc, rc) = self.chroma_at(self.last_qidx);
        (qs, self.qstep_base as f32 / qs as f32, qc, rc)
    }

    /// Serial raster-order pre-pass producing the per-SB AQ grid, **bit-exact** with
    /// the serial `per_sb`(full-interior SBs) / `current`(edge & `!present`) sequence
    /// the main loop runs. This removes the raster-serial `last_qidx` dependency so a
    /// wavefront (diagonal) decide can read `grid[r*sb_cols + c]` instead of calling
    /// `per_sb`; the serial emit pass then replays `sig` into `enc.delta_q_signaled`.
    /// `per_sb` is invoked ⟺ `full_interior` (whole-64 fast path OR `sb_use_split`);
    /// edge SBs (`sb_walk`) use `current()` — mirrored here exactly.
    #[allow(dead_code)]
    pub(crate) fn precompute_grid(
        &self,
        yp: &[f32],
        pw: usize,
        width: usize,
        height: usize,
        needs_partition: bool,
    ) -> Vec<AqCell> {
        let sb_cols = width.div_ceil(64);
        let sb_rows = height.div_ceil(64);
        let step = 1i32 << AQ_RES_LOG2;
        let mut last_qidx = self.last_qidx;
        let mut grid = Vec::with_capacity(sb_cols * sb_rows);
        for r in 0..sb_rows {
            for c in 0..sb_cols {
                let sb_y = r * 64;
                let sb_x = c * 64;
                // `per_sb` (accumulate) is called ⟺ the SB is NOT a native edge walk,
                // i.e. `!sb_walk` = `full_interior || !needs_partition`. When the frame
                // is padded (`!needs_partition`) even geometric-edge SBs take the
                // whole-64 path and accumulate; only native-partition edge SBs
                // (`needs_partition && !full_interior`) use `current()` (no accumulate).
                let full_interior = (sb_x + 64 <= width && sb_y + 64 <= height) || !needs_partition;
                // Accumulator qstep entering this SB (== `aqs.current()` at the probe).
                let qs_before = if self.present {
                    qstep(last_qidx as u32) as i32
                } else {
                    self.qstep_base
                };
                let resid_scale_before = self.qstep_base as f32 / qs_before as f32;
                let cell = if !self.present {
                    let (qc, rc) = self.chroma_at(self.base_q);
                    AqCell {
                        qs: self.qstep_base,
                        resid_scale: 1.0,
                        qs_c: qc,
                        resid_scale_c: rc,
                        sig: 0,
                        qidx: self.base_q,
                        qs_before,
                        resid_scale_before,
                    }
                } else if full_interior {
                    // Mirrors `per_sb`: variance → target → signalled delta → newq (accumulate).
                    let target = self.sb_target(yp, pw, sb_y, sb_x, width, height);
                    let sig = (((target - last_qidx) as f32) / step as f32)
                        .fast_round()
                        .clamp(-(AQ_MAX_SIGNALED as f32), AQ_MAX_SIGNALED as f32)
                        as i32;
                    let newq = (last_qidx + sig * step).clamp(1, 255);
                    last_qidx = newq;
                    let qs = qstep(newq as u32) as i32;
                    let (qc, rc) = self.chroma_at(newq);
                    AqCell {
                        qs,
                        resid_scale: self.qstep_base as f32 / qs as f32,
                        qs_c: qc,
                        resid_scale_c: rc,
                        sig,
                        qidx: newq,
                        qs_before,
                        resid_scale_before,
                    }
                } else {
                    // Mirrors `current()`: no accumulation, no signalled delta.
                    let qs = qstep(last_qidx as u32) as i32;
                    let (qc, rc) = self.chroma_at(last_qidx);
                    AqCell {
                        qs,
                        resid_scale: self.qstep_base as f32 / qs as f32,
                        qs_c: qc,
                        resid_scale_c: rc,
                        sig: 0,
                        qidx: last_qidx,
                        qs_before,
                        resid_scale_before,
                    }
                };
                grid.push(cell);
            }
        }
        grid
    }

    #[allow(clippy::too_many_arguments)]
    pub(crate) fn per_sb(
        &mut self,
        enc: &mut RangeEncoder,
        yp: &[f32],
        pw: usize,
        sb_y: usize,
        sb_x: usize,
        width: usize,
        height: usize,
    ) -> (i32, f32, i32, f32) {
        if !self.present {
            enc.delta_q_signaled = 0;
            let (qc, rc) = self.chroma_at(self.base_q);
            return (self.qstep_base, 1.0, qc, rc);
        }
        // Variance Boost (flat-region protection) combined with dark-structured-detail
        // protection, signalled through the same accumulator/qstep machinery as classic AQ.
        let target = self.sb_target(yp, pw, sb_y, sb_x, width, height);
        let step = 1i32 << AQ_RES_LOG2;
        let sig = (((target - self.last_qidx) as f32) / step as f32)
            .fast_round()
            .clamp(-(AQ_MAX_SIGNALED as f32), AQ_MAX_SIGNALED as f32) as i32;
        let newq = (self.last_qidx + sig * step).clamp(1, 255);
        self.last_qidx = newq;
        enc.delta_q_signaled = sig;
        let qs = qstep(newq as u32) as i32;
        let (qc, rc) = self.chroma_at(newq);
        (qs, self.qstep_base as f32 / qs as f32, qc, rc)
    }
}

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

    #[test]
    fn dark_protection_targets_dark_structure_only() {
        // 64x64 SBs at three luma levels, each with a coherent striped texture (survives
        // the 2x downsample, so it registers as structure not noise).
        let build = |mean: f32| -> Vec<f32> {
            let mut p = vec![0f32; 64 * 64];
            for r in 0..64 {
                for c in 0..64 {
                    // 4px-wide stripes: strong at both scales.
                    p[r * 64 + c] = mean + if (c / 4) % 2 == 0 { 12.0 } else { -12.0 };
                }
            }
            p
        };
        let base_q = 190; // in the gated range (>= 150)
        let st =
            AqState::new(true, base_q, qstep(base_q as u32) as i32, 10.0, 0).with_dark_aq(DarkAq {
                enabled: true,
                ..DarkAq::default()
            });
        let dark = build(36.0);
        let bright = build(180.0);
        let flat_dark = vec![24.0f32; 64 * 64]; // dark but no structure
        let d_dark = st.dark_protection(&dark, 64, 0, 0, 64, 64);
        let d_bright = st.dark_protection(&bright, 64, 0, 0, 64, 64);
        let d_flat = st.dark_protection(&flat_dark, 64, 0, 0, 64, 64);
        assert!(
            d_dark > 0,
            "dark structured SB should be protected, got {d_dark}"
        );
        assert_eq!(d_bright, 0, "bright structured SB must not be protected");
        assert_eq!(
            d_flat, 0,
            "dark flat SB (no structure) must not be protected"
        );
        // Gate: disabled below min_q.
        let st_hi = AqState::new(true, 100, qstep(100) as i32, 10.0, 0).with_dark_aq(DarkAq {
            enabled: true,
            ..DarkAq::default()
        });
        assert_eq!(
            st_hi.dark_protection(&dark, 64, 0, 0, 64, 64),
            0,
            "gated out below min_q"
        );
    }

    #[test]
    fn precompute_grid_matches_serial_per_sb() {
        // The wavefront AQ pre-pass must be bit-exact with the serial
        // `per_sb`(full-interior) / `current`(edge) sequence the main loop runs.
        let base_q = 120;
        let base_step = qstep(base_q as u32) as i32;
        // Non-SB-aligned dims so the grid mixes full-interior AND edge SBs.
        let (width, height) = (160usize, 96usize);
        let pw = crate::av2::helpers::sb_align(width);
        let ph = crate::av2::helpers::sb_align(height);
        let yp: Vec<f32> = (0..pw * ph).map(|i| (i * 37 % 251) as f32).collect();
        let mk = || AqState::new(true, base_q, base_step, 10.0, 0);
        let sb_cols = width.div_ceil(64);
        let sb_rows = height.div_ceil(64);

        // `needs_partition` selects the edge rule: `true` = native edge SBs use
        // `current()` (no accumulate); `false` = padded frame, edge SBs also
        // `per_sb`. The main loop's accumulate condition is `!sb_walk`
        // (`full_interior || !needs_partition`) — mirror it in the reference.
        for needs_partition in [true, false] {
            let grid = mk().precompute_grid(&yp, pw, width, height, needs_partition);
            let mut aq = mk();
            let mut enc = crate::av2::entropy::RangeEncoder::new();
            for r in 0..sb_rows {
                for c in 0..sb_cols {
                    let (sb_y, sb_x) = (r * 64, c * 64);
                    let full_interior =
                        (sb_x + 64 <= width && sb_y + 64 <= height) || !needs_partition;
                    let (qs, rs, qc, rc) = if full_interior {
                        aq.per_sb(&mut enc, &yp, pw, sb_y, sb_x, width, height)
                    } else {
                        enc.delta_q_signaled = 0;
                        aq.current()
                    };
                    let cell = grid[r * sb_cols + c];
                    assert_eq!(cell.qs, qs, "qs at ({r},{c}) np={needs_partition}");
                    assert_eq!(cell.qs_c, qc, "qs_c at ({r},{c}) np={needs_partition}");
                    assert_eq!(cell.sig, enc.delta_q_signaled, "sig at ({r},{c})");
                    assert_eq!(cell.resid_scale.to_bits(), rs.to_bits(), "rs at ({r},{c})");
                    assert_eq!(cell.resid_scale_c.to_bits(), rc.to_bits(), "rs_c ({r},{c})");
                    assert_eq!(cell.qidx, aq.current_qidx(), "qidx at ({r},{c})");
                }
            }
        }
    }

    #[test]
    fn probe_does_not_advance_delta_q_accumulator() {
        let base_q = 120;
        let base_step = qstep(base_q as u32) as i32;
        let mut aq =
            AqState::new(true, base_q, base_step, 10.0, 0).with_variance_boost(6, 1.0, true);
        let plane = vec![128.0f32; 64 * 64];

        let before = aq.current();
        let probed = aq.per_sb_probe(&plane, 64, 0, 0, 64, 64);
        assert_ne!(probed.0, before.0, "flat block should request an AQ boost");
        assert_eq!(
            aq.current().0,
            before.0,
            "probe must preserve decoder qindex state"
        );

        let mut enc = RangeEncoder::new();
        enc.delta_q_present = true;
        let committed = aq.per_sb(&mut enc, &plane, 64, 0, 0, 64, 64);
        assert_eq!(committed.0, probed.0);
        assert_eq!(aq.current().0, probed.0);
        assert_ne!(enc.delta_q_signaled, 0);
    }
}